EntitySubstituteTest.py

#!/usr/bin/env python
# coding: utf-8

import json
import os.path

from typing import Dict, Callable
from itertools import permutations
from transformers import AutoModelForMaskedLM, AutoTokenizer, BertForMaskedLM
import torch
from torch import nn
from tqdm import tqdm
import copy
import pickle
from torch.nn.functional import pad
from collections import defaultdict
from pathlib import Path
import time
import requests
torch.cuda.empty_cache()
from fitbert import *
from scoringmethod import *
from document import *

BATCH_SIZE = 512  # BS number, don't remember what it was, isn't used in the experiments. Remove after refactor.
    
# FitBert.nonlins = {None: lambda x:x,
#                "softmax": FitBert.softmax,
#                "relu": torch.relu,
#                "relmax": FitBert.masked_softmax,
#                "top10": FitBert.soft_top_k,
#                "top20": lambda x: FitBert.soft_top_k(x, 20),
#                "top50": lambda x: FitBert.soft_top_k(x, 50),
#                "top100": lambda x: FitBert.soft_top_k(x, 100)
#               }

def read_punctuation(fb):
    punctuation = fb.tokenizer([".,-()[]{}_=+?!@#$%^&*\\/\"'`~;:|…）（•−"], add_special_tokens=False)['input_ids'][0]
    one_hot_punctuation = torch.ones(fb.bert.get_output_embeddings().out_features, dtype=torch.long)
    one_hot_punctuation[punctuation] = 0
    one_hot_punctuation[1232] = 0
    return one_hot_punctuation


def candidates(prompt:str, choices, return_ments=False):
    for a, b in permutations(choices, 2):
        if return_ments:
            yield prompt.replace("?x", a, 1).replace("?y", b, 1), (a, b)
        else:
            yield prompt.replace("?x", a, 1).replace("?y", b, 1)


# scores equivalently to the old method, even with padding.
# Can be used to batch across examples.
def pll_score_batched(self, sents: list, return_all=False):
    self.bert.eval()
    key_to_sent = {}
    with torch.no_grad():
        data = {}
        for sent in sents:
            tkns = self.tokenizer.tokenize(sent)
            data[len(data)] = {
                'tokens': tkns,
                'len': len(tkns)
            }
        scores = {"pll": {}}
        all_plls = {"pll": {}}

        sents_sorted = list(sorted(data.keys(), key=lambda k: data[k]['len']))

        inds = []
        lens = []

        methods = [PLL()]

        for sent in sents_sorted:
            n_tokens = data[sent]['len']
            if sum(lens) <= BATCH_SIZE:
                inds.append(sent)
                lens.append(n_tokens)
            else:
                # There is at least one sentence.
                # If the count is zero, then its size is larger than the batch size.
                # Send it anyway.
                flag = (len(inds) == 0)
                if flag:
                    inds.append(sent)
                    lens.append(n_tokens)
                _inner_score_stuff(self, data, inds, lens, methods, scores, all_plls, return_all)
                inds = [sent]
                lens = [n_tokens]
            if sent == sents_sorted[-1]:
                _inner_score_stuff(self, data, inds, lens, methods, scores, all_plls, return_all)

        for d in data:
            data[d].clear()
        data.clear()
        # del all_probs
        if self.device == "cuda":
            torch.cuda.empty_cache()
        for method in scores:
            assert len(scores[method]) == len(sents_sorted)
        if return_all:
            return unsort_flatten(scores)["pll"], unsort_flatten(all_plls)["pll"]
        return unsort_flatten(scores)["pll"]


def unsort_flatten(mapping):
    # print(mapping.keys())
    return {f: list(mapping[f][k] for k in range(len(mapping[f]))) for f in mapping}


def cos_score_batched(self, sents: list, return_all=True):
    return score_batched(self, methods=[CSD()], sents=sents, return_all=return_all)


def euc_score_batched(self, sents: list, return_all=True):
    return score_batched(self, methods=[ESD()], sents=sents, return_all=return_all)


def jsd_score_batched(self, sents: list, return_all=True):
    return score_batched(self, methods=[JSD()], sents=sents, return_all=return_all)


def msd_score_batched(self, sents: list, return_all=True):
    return score_batched(self, methods=[MSD()], sents=sents, return_all=return_all)


def hel_score_batched(self, sents: list, return_all=True):
    return score_batched(self, methods=[HSD()], sents=sents, return_all=return_all)


def all_score_batched(self, sents: list,  return_all=True):
    return score_batched(self, methods=list(KNOWN_METHODS.values()), sents=sents, return_all=return_all)


# # Purely for reference.
# def mask_tokenize(self, sent, keep_original=False, add_special_tokens=False, padding=False, return_full=False):
#     print("S:", type(sent), sent)
#     tokens = self.tokenize(sent, add_special_tokens=add_special_tokens, padding=padding)
#     # print(tokens)
#     tlen = len(tokens)
#     offset = 1 if add_special_tokens else 0
#     token_mat = [tokens[:] for i in range(tlen - (2*offset))]
#     for i in range(offset, tlen-offset):
#         token_mat[i-offset][i] = self.tokenizer.mask_token
#     if keep_original:
#         token_mat = [tokens[:]] + token_mat

#     if return_full:
#         return token_mat, self.tokenizer(token_mat, add_special_tokens=(not add_special_tokens), is_split_into_words=True, return_tensors='pt')
#     return token_mat


# For this one:
# Take a sentence, tokenize it, return all needed information like number of tokens.
def _inner_tokenize_sentence(self, sent, keep_original):
    _, tkns = self.mask_tokenize(sent, keep_original=keep_original, add_special_tokens=True, return_full=True)
    # print(tkns)
    # print(f"TKNS:{len(tkns.input_ids[0]) - 2}")
    return tkns, len(tkns.input_ids[0]) - 2


def score_batched(self, methods, sents: list, return_all=True):
    # Enforce evaluation mode
    self.bert.eval()
    with torch.no_grad():
        data = {}
        for sent in sents:
            # Tokenize every sentence
            # print("S2:", sent)
            tkns, n_tkns = _inner_tokenize_sentence(self, sent, keep_original=True)
            data[len(data)] = {
                'tokens': tkns,
                'len': n_tkns
            }
        # print("Boo")

        scores = {m.label: {} for m in methods}
        all_plls = {m.label: {} for m in methods}

        sents_sorted = list(sorted(data.keys(), key=lambda k: data[k]['len']))

        inds = []
        lens = []

        for sent in sents_sorted:
            n_tokens = data[sent]['len']
            if sum(lens) <= BATCH_SIZE:
                inds.append(sent)
                lens.append(n_tokens)
            else:
                # There is at least one sentence.
                # If the count is zero, then its size is larger than the batch size.
                # Send it anyway.
                flag = (len(inds) == 0)
                if flag:
                    inds.append(sent)
                    lens.append(n_tokens)
                _inner_score_stuff(self, data, inds, lens, methods, scores, all_plls, return_all)
                inds = [sent]
                lens = [n_tokens]
            if sent == sents_sorted[-1]:
                _inner_score_stuff(self, data, inds, lens, methods, scores, all_plls, return_all)

        for d in data:
            data[d].clear()
        data.clear()
        if self.device == "cuda":
            torch.cuda.empty_cache()
        for method in scores:
            assert len(scores[method]) == len(sents_sorted)
        if return_all:
            return unsort_flatten(scores), unsort_flatten(all_plls)
        return unsort_flatten(scores)


def bert_am(self, data, *args, **kwds):
    return self.bert(data, *args, attention_mask=(data!=self.tokenizer.pad_token_id), **kwds)

def _inner_score_stuff(self, data, inds, lens, methods, scores, all_plls, return_all):
    longest = max(lens)

    bert_forward = torch.concat([pad(data[d]['tokens'].input_ids, (0, longest - l), 'constant', self.tokenizer.pad_token_id ) for d,l in zip(inds, lens)], dim=0).to(self.device)
    token_type_ids = torch.concat([pad(data[d]['tokens'].token_type_ids, (0, longest - l), 'constant', 0) for d, l in zip(inds, lens)], dim=0).to(self.device)
    _probs = self.softmax(bert_am(self, bert_forward, token_type_ids=token_type_ids)[0])[:, 1:, :]
    
    del bert_forward

    use_pll = any(["pll" in method.label for method in methods])
    print(["pll" in method.label for method in methods])
    print(use_pll)

    for ind, slen in zip(inds, lens):
        origids = data[ind]['tokens'].input_ids[0][1:-1].to(self.device) if use_pll else None
        for method in methods:
            prob, alls = method(_probs[:slen + 1], origids=origids, return_all=True)
            if return_all:
                assert ind not in all_plls[method.label]
                all_plls[method.label][ind] = alls
            assert ind not in scores[method.label]
            scores[method.label][ind] = prob
            del alls, prob
        _probs = _probs[slen + 1:]
    del _probs


def extend_bert(fb: FitBert, num_blanks:int, tokens_per_blank:int) -> FitBert:
    expandable_tokens = ['?is', '?the', '?The', '?award', '?sibling', '?team', '?spouse']
    add_tokens = ['?x', '?y'] + expandable_tokens
    if fb.uses_bpe:
        add_tokens.extend([f" {t}" for t in add_tokens])
        add_tokens.append(f' {fb.mask_token}')
    for e in range(num_blanks):
        for t in range(tokens_per_blank):
            add_tokens.append(f"[ENT_{e}_{t}]")
            if fb.uses_bpe:
                add_tokens.append(f" [ENT_{e}_{t}]")
        add_tokens.append(f"[ENT_{e}_x]")  # Quick hack to help with preprocessing
        if fb.uses_bpe:
            add_tokens.append(f" [ENT_{e}_x]")
    add_tokens.append('[ENT_BEG]')
    add_tokens.append('[ENT_END]')
    if fb.uses_bpe:
        add_tokens.append(' [ENT_BEG]')
        add_tokens.append(' [ENT_END]')
    fb.tokenizer.add_tokens(add_tokens, special_tokens=True)  # Add the tokens to the tokenizer.
    fb.bert.resize_token_embeddings(len(fb.tokenizer))  # Add the tokens to the embedding matrix, initialize with defaults. DO NOT TRAIN.
    fb.token_width = tokens_per_blank
    fb.entity_tokens = fb.tokenizer("".join(add_tokens), add_special_tokens=False)['input_ids']
    fb.input_embeddings = fb.bert.get_input_embeddings().to(dtype=fb.float_dtype)
    return fb


def repair_masking(masked_doc):
    masked = ""
    for token in masked_doc:
        if token[:2] == "##":
            masked += token[2:]
        else:
            if masked:
                masked += " "
            masked += token
    return masked


def mask_vectors(self, sent, keep_original=False, add_special_tokens=False, pad_to=0):
    # tokens = self.tokenize(sent, add_special_tokens=add_special_tokens, padding=padding)
    # print(tokens)
    sent.squeeze_(0)
    # print(sent.shape)
    tlen = len(sent)
    offset = 1 if add_special_tokens else 0
    if pad_to > 0:
        pad_length = pad_to - tlen
        if pad_length > 0:
            sent = torch.cat( (sent, torch.clone(self.pad_token_vector).repeat((pad_length, 1)).to(sent.device)), dim=-2)
    else:
        pad_to = tlen
    token_mat = [torch.clone(sent) for i in range(tlen - (2*offset))]
    for i in range(offset, tlen-offset):
        # print("ti B:", token_mat[i-offset][i])
        token_mat[i-offset][i] = self.mask_token_vector
        # print("ti A:", token_mat[i-offset][i])
    if keep_original:
        token_mat = [torch.clone(sent)] + token_mat
    
    mask = torch.ones((len(token_mat), pad_to), dtype=torch.int)
    mask[:, tlen:pad_to] = 0
    # print(pad_to, tlen, len(token_mat), mask.shape)
    assert mask.shape[-1] == pad_to
    return torch.stack(token_mat), mask


def score_vectors(self, methods, probs, return_all=True):
    # Enforce evaluation mode
    self.bert.eval()
    with torch.no_grad():
        use_pll = "pll" in [method.label for method in methods]

        assert not use_pll, "PLL cannot be used for raw vectors. Please choose a different metric."

        all_plls = {}
        scores = {}
        
        for method in methods:
            prob, alls = method(probs, return_all=True)
            if return_all:
                all_plls[method.label] = alls
            scores[method.label] = prob

        if self.device == "cuda":
            torch.cuda.empty_cache()
        # for method in scores:
        #     assert len(scores[method]) == len(sents_sorted)
        if return_all:
            return scores, all_plls
        return scores
    
jfile = None
def load_data_json(path:str, task_name:str, dset:str):
    global jfile
    if not jfile:
        if ('docred' in task_name or 're-docshred' in task_name) and dset == 'train':
            dset = 'train_annotated'
        with open(f"{path}/{task_name}/{dset}.json") as datafile:
            jfile = json.load(datafile)
    return jfile

def process_single_document(datadir, resdir, fb:FitBert, task_name='docred', dset='dev', doc=0, num_blanks=2, num_passes=1, use_ent=False, top_k=0):
    with torch.no_grad():
        jfile = load_data_json(datadir, task_name, dset)
        if doc >= len(jfile):
            raise StopIteration
        d = Document(jfile[doc], doc, width=num_blanks, mlm=fb, num_passes=num_passes, use_ent=use_ent)
        docfile = f"{resdir}/{task_name}_{fb.model_name.replace('/', '_')}_{dset}_{doc}_{num_blanks}b_{num_passes}p.pickle"
        # docfile = f"{resdir}/{(task_name + '_') if task_name != 'docred' else ''}{(model + '_') if model != 'bert-large-cased' else ''}{dset}_{d.num}{'_' + str(num_blanks) + 'blanks' if num_blanks != 2 else ''}{'' if use_ent else '_MASK'}{'' if num_passes == 1 else '_'+str(num_passes)}.pickle"
        # print(docfile)
        # if os.path.exists(docfile) or
        if d.num in skip or d.num < start_at:
            # print(f"Pass {p_doc.num}")
            # print(f"Document {d.num} skipped.", flush=True)
            return None, None
        # print(f"Document {d.num} started.", flush=True)
        # print(d.doc["title"])
        md = d.masked_doc
        # print(len(md['tokens']))
        if len(md['tokens']) <= (fb.context_window-2):
            mask = [False] + md['ment_mask'] + [False]
            # print("mt", len(md['tokens']))
            # ents = md['ents']
            _tokens = fb.tokenizer.convert_tokens_to_ids(md['tokens'])
            _tokens = [fb.tokenizer.cls_token_id] + _tokens + [fb.tokenizer.sep_token_id]
            _tokens = torch.LongTensor([_tokens])
            V = fb.get_vocab_output_dim()  # [tokens, vocab(29028)]

            d.ment_inds = _tokens.squeeze(0)[mask]
            # d.ment_inds_masked = d.ment_inds.clone()
            # d.ment_inds_masked[d.ment_inds >= min(fb.entity_tokens)] = -1
            out: Dict[int, torch.Tensor] = dict()
            # Initial pass: Just one-hot vectors as inputs.
            # print("t", _tokens.shape)
            input_vecs = torch.nn.functional.one_hot(_tokens, V).to(fb.float_dtype)
            out[0] = input_vecs.cpu()
            # if num_passes == 0:
            #     # Then take the input tokens and convert them to one-hot vectors.
            # else:
            #     input_vecs = fb.input_embeddings(_tokens.to(fb.device)).cpu()
            #     out[0] = fb.bert(inputs_embeds=input_vecs.to(fb.device)).logits.cpu()
            #     d.ment_inds = torch.LongTensor([-1]*sum(mask))
            # del _tokens
            np = num_passes
            cp = 1
            while np > 0:
                # Second and further passes: Run through the MLM.
                # Step 1: Take original input vecs and sub in the entities.
                if cp > 1:
                    entities = out[cp - 1].squeeze(0)[mask]
                    if top_k > 0:
                        entities = fb.fuzzy_embed(fb.soft_top_k(entities, k=top_k))
                    elif top_k == 0:
                        entities = fb.fuzzy_embed(fb.softmax_(entities.clone()))
                    else:
                        entities = fb.fuzzy_embed(entities)
                    input_embeds.squeeze(0)[mask] = entities.squeeze(0)
                else:
                    # First pass: Just use normal input embeddings.
                    input_embeds = fb.input_embeddings(_tokens.to(fb.device))

                # print(f"{cp}: {input_embeds.dtype}!!")
                # else it's the same vectors already.
                # Then we do a forward pass, gather the new output logits.
                with torch.autocast("cuda", dtype=fb.float_dtype):
                    out[cp] = fb.bert(inputs_embeds=input_embeds).logits.detach().float().cpu()
                # If bert is working with lower precision, increase it here.
                # if reduced_precision:
                #     # for cp in out:
                #     out[cp].float()
                # print(f"{cp}: {out[cp].dtype}!!")
                # print(cp, out[cp].shape)
                cp += 1
                np -= 1
            del _tokens
            
            d.ment_vecs = dict()
            for p in out:
                if num_blanks > 0:
                    d.ment_vecs[p] = out[p].squeeze(0)[mask].view(-1, num_blanks, V)
                else:
                    d.ment_vecs[p] = out[p].squeeze(0)[mask].view(-1, V)
                # print(d.ment_vecs[p].shape)
            del out
            torch.cuda.empty_cache()
            # print(f"Document {d.num} preprocessed.", flush=True)
            return d, docfile
            # exit(0)
        else:
            print("skipped", d.num)
            return None, None

def process_documents(resdir, fb:FitBert = None, task_name='docred', dset='dev', doc=0, num_blanks=2, num_passes=1, use_ent=False, top_k=0, skip=[], model='bert-large-cased', start_at=0):
    with torch.no_grad():  # Super enforce no gradients whatsoever.
        torch.cuda.empty_cache()
        if fb is None:
            fb = extend_bert(FitBert(model_name=model), 200, num_blanks)
            # nls = {None: lambda x:x, "softmax": FitBert.softmax, "relu":torch.relu}
            global one_hot_punctuation
            one_hot_punctuation = read_punctuation(fb)
        # qx, qy = fb.tokenizer(["?x?y"], add_special_tokens=False)['input_ids'][0]
        # Take a document
        # Find all the entities
        processed_docs = []
        tot = 1000 if dset == 'dev' else 3053
        d: Document = ...  # Silly way to get some IDEs to give better completions.
        for d in read_document(task_name=task_name, dset=dset, doc=doc, num_blanks=num_blanks, mlm=fb, path='data', use_ent=use_ent):
            docfile = f"{resdir}/{task_name}_{model}_{dset}_{d.num}_{num_blanks}b_{num_passes}p.pickle"
            # docfile = f"{resdir}/{(task_name + '_') if task_name != 'docred' else ''}{(model + '_') if model != 'bert-large-cased' else ''}{dset}_{d.num}{'_' + str(num_blanks) + 'blanks' if num_blanks != 2 else ''}{'' if use_ent else '_MASK'}{'' if num_passes == 1 else '_'+str(num_passes)}.pickle"
            # print(docfile)
            # if os.path.exists(docfile) or
            if d.num in skip or d.num < start_at:
                # print(f"Pass {p_doc.num}")
                # print(f"Document {d.num} skipped.", flush=True)
                continue
            print(f"Document {d.num} started.", flush=True)
            print(d.doc["title"])
            md = d.masked_doc
            print(len(md['tokens']))
            if len(md['tokens']) <= (fb.context_window-2):
                mask = [False] + md['ment_mask'] + [False]
                # print("mt", len(md['tokens']))
                # ents = md['ents']
                _tokens = fb.tokenizer.convert_tokens_to_ids(md['tokens'])
                _tokens = [fb.tokenizer.cls_token_id] + _tokens + [fb.tokenizer.sep_token_id]
                _tokens = torch.LongTensor([_tokens])
                V = fb.get_vocab_output_dim()  # [tokens, vocab(29028)]

                d.ment_inds = _tokens.squeeze(0)[mask]
                # d.ment_inds_masked = d.ment_inds.clone()
                # d.ment_inds_masked[d.ment_inds >= min(fb.entity_tokens)] = -1
                out: Dict[int, torch.Tensor] = dict()
                # Initial pass: Just one-hot vectors as inputs.
                # print("t", _tokens.shape)
                input_vecs = torch.nn.functional.one_hot(_tokens, V).to(fb.float_dtype)
                out[0] = input_vecs.cpu()
                # if num_passes == 0:
                #     # Then take the input tokens and convert them to one-hot vectors.
                # else:
                #     input_vecs = fb.input_embeddings(_tokens.to(fb.device)).cpu()
                #     out[0] = fb.bert(inputs_embeds=input_vecs.to(fb.device)).logits.cpu()
                #     d.ment_inds = torch.LongTensor([-1]*sum(mask))
                # del _tokens
                np = num_passes
                cp = 1
                while np > 0:
                    # Second and further passes: Run through the MLM.
                    # Step 1: Take original input vecs and sub in the entities.
                    if cp > 1:
                        entities = out[cp - 1].squeeze(0)[mask]
                        if top_k > 0:
                            entities = fb.fuzzy_embed(fb.soft_top_k(entities, k=top_k))
                        elif top_k == 0:
                            entities = fb.fuzzy_embed(fb.softmax_(entities.clone()))
                        else:
                            entities = fb.fuzzy_embed(entities)
                        input_embeds.squeeze(0)[mask] = entities.squeeze(0)
                    else:
                        # First pass: Just use normal input embeddings.
                        input_embeds = fb.input_embeddings(_tokens.to(fb.device))

                    # print(f"{cp}: {input_embeds.dtype}!!")
                    # else it's the same vectors already.
                    # Then we do a forward pass, gather the new output logits.
                    with torch.autocast("cuda", dtype=fb.float_dtype):
                        out[cp] = fb.bert(inputs_embeds=input_embeds).logits.detach().float().cpu()
                    # If bert is working with lower precision, increase it here.
                    # if reduced_precision:
                    #     # for cp in out:
                    #     out[cp].float()
                    # print(f"{cp}: {out[cp].dtype}!!")
                    print(cp, out[cp].shape)
                    cp += 1
                    np -= 1
                del _tokens
                
                d.ment_vecs = dict()
                for p in out:
                    if fb.token_width > 0:
                        d.ment_vecs[p] = out[p].squeeze(0)[mask].view(-1, fb.token_width, V)
                    else:
                        d.ment_vecs[p] = out[p].squeeze(0)[mask].view(-1, V)
                    # print(d.ment_vecs[p].shape)
                del out
                torch.cuda.empty_cache()
                print(f"Document {d.num} preprocessed.", flush=True)
                yield d, docfile
                # exit(0)
            else:
                print("skipped", d.num)


def replace_embeddings(x, y, prompt):
    ix = prompt['ix']
    iy = prompt['iy']
    embs = prompt['vecs'].clone()
    # tkns = prompt['input_ids']
    return torch.cat([embs[:,:ix], x, embs[:,ix+1:iy], y, embs[:,iy+1:]], dim=1)


def replace_ids(x, y, prompt):
    ix = prompt['ix']
    iy = prompt['iy']
    embs = prompt['input_ids']
    return torch.cat([embs[:ix], x, embs[ix+1:iy], y, embs[iy+1:]])


def output_to_fuzzy_embeddings(fb: FitBert, v: torch.Tensor):
    # print(v.to(device=fb.device, dtype=fb.float_dtype)@fb.input_embeddings.weight)
    # print("=================!")
    # return vec.to(device=self.device, dtype=self.float_dtype)@self.input_embeddings.weight
    return (v.to(device=fb.device, dtype=fb.float_dtype)@fb.input_embeddings.weight).cpu()


def meminfo():
    f, t = torch.cuda.mem_get_info()
    f = f / (1024 ** 3)
    t = t / (1024 ** 3)
    return f"{f:.2f}g/{t:.2f}g"


expansion_rules = {
    '?is': ['is', 'was'],
    '?the': ['the', ''],
    '?The': ['The', ''],
    '?award': ['award', ''],
    '?sibling': ['sibling', 'sister', 'brother'],
    '?team': ['team', ''],
    '?spouse': ['spouse', 'husband', 'wife']
}

from typing import List
def expand_prompt(prompt:str, i:int=0) -> List[str]:
    # print(i, prompt)
    prompts = set()
    for rule, expansion in expansion_rules.items():
        if rule in prompt:
            for exp in expansion:
                prompts.update(expand_prompt(prompt.replace(rule, exp, 1).replace('  ', ' ').replace(' .', '.').lstrip(), i + 1))
            return list(prompts)
    return [prompt]


def get_prompt_info(prompts, rel_info, fb: FitBert):
    qx, qy = fb.tokenizer(["?x?y"], add_special_tokens=False)['input_ids'][0]
    if fb.uses_bpe:
        sqx, sqy = fb.tokenizer([" ?x ?y"], add_special_tokens=False)['input_ids'][0]
    prompt_data = {}
    for prompt in prompts:
        prompt_data[prompt] = list()
        for ep in expand_prompt(rel_info[prompt]['prompt_xy']):
            pi = dict()
            tkns = fb.tokenizer(ep, return_tensors='pt')['input_ids']    
            if fb.uses_bpe:
                tkns[tkns == sqx] = qx
                tkns[tkns == sqy] = qy
            pi['input_ids'] = tkns[0].cpu()
            pi['ix'] = torch.where(tkns[0] == qx)[0].item()
            pi['iy'] = torch.where(tkns[0] == qy)[0].item()
            # print("Cosine score:", score_batched(fb, [scorer], [prompt])[0]['csd'][0])
            pi['vecs'] = fb.input_embeddings(tkns.to(device=fb.device)).cpu()
            prompt_data[prompt].append(pi)
    # print(prompt_data)
    return prompt_data


def run_many_experiments_orig(task_name, dset, rel_info, nonlins, poolers, scorers, resdir, num_blanks, num_passes=1, max_batch=2000, use_ent=False, skip=[], stopfile='', model='bert-large-cased', start_at=0, reduced_precision=False, big_batch=False):
    # import time
    if num_blanks == 0:
        poolers = [None]
    with torch.no_grad():
        torch.cuda.empty_cache()
        fb = extend_bert(FitBert(model_name=model, reduced_precision=reduced_precision), 200, num_blanks)
        model = model.split('/')[-1]
        # nls = {None: lambda x:x, "softmax": FitBert.softmax, "relu":torch.relu}
        # global one_hot_punctuation
        # one_hot_punctuation = read_punctuation(fb)
        if task_name == "docred" or task_name == "docshred":
            prompts = ['P17', 'P27', 'P131', 'P150', 'P161', 'P175', 'P527', 'P569', 'P570', 'P577']
        elif task_name == "biored":
            prompts = ['Association', 'Bind', 'Negative_Correlation', 'Positive_Correlation']
        elif task_name == "dwie":
            prompts = ['gpe', 'in', 'citizen_of', 'based_in', 'member_of', 'agent_of', 'head_of', 'head_of_state', 'agency_of', 'head_of_gov']  # ['minister_of', 'event_in', 'part_of', 'institution_of', 'appears_in', 'award_received', 'ministry_of', 'player_of', 'vs', 'created_by']
        else:
            prompts = list(sorted(rel_info.keys()))
        for p in prompts:
            assert p in rel_info, f"{p} not found in rel_info. Is it spelled wrong?"
        # prompts = ['P17']
        prompt_data = get_prompt_info(prompts, rel_info, fb)

        otime = 0
        ntime = 0
        for p_doc, docfile in process_documents(resdir, fb=fb, task_name=task_name, dset=dset, doc=-1, num_blanks=num_blanks, num_passes=num_passes, use_ent=use_ent, skip=skip, model=model, start_at=start_at, top_k=0):
            all_scores = dict()
            for nps in range(0, num_passes):
                print(f"{nps=}")
                if num_blanks > 0 and nps == 0:
                    continue
                if os.path.exists(docfile.replace(f'_{num_passes}p', f'_{nps}p')):
                    print(f"Document {p_doc.num} at {nps} passes skipped.", flush=True)
                    continue
                all_scores = dict()
                for nonlin in nonlins:
                    print(f"{nonlin=}")
                    # print(f"NL: {nonlin}")
                    all_scores[nonlin] = {}
                    # nl = FitBert.nonlins[nonlin]
                    for pooler in poolers:
                        print(f"{pooler=}")
                        # print(f"PL: {pooler}")
                        all_scores[nonlin][pooler] = {}
                        evs, ev_tkns = p_doc.entity_vecs(nonlinearity=nonlin, pooling=pooler, passes=nps)
                        fuzzy_embeds = {e:[output_to_fuzzy_embeddings(fb, v1.unsqueeze(0)) for v1 in evs[e]] for e in evs}
                        e_to_m_map = p_doc.masked_doc['ents']
                        # print(f"A: {torch.cuda.mem_get_info()}")
                        # times = []
                        for prompt_id in prompts: # rel_info:
                            # print(f"{prompt_id=}")
                            if prompt_id.split('|')[0] not in p_doc.relations:
                                continue
                            # nnow = time.time()
                            # print(f"PI: {prompt_id}")
                            ans = [(a[0], a[2]) for a in p_doc.answers(detailed=False) if a[1] == prompt_id.split('|')[0]]
                            # BioRED explicitly states that all relations are non-directional.
                            # This is honestly false, but we marked the ones that aren't clearly non-directional to avoid issues.
                            # For example, "Conversion" is a one-way process between chemicals.
                            # "Bind" is questionable in this regard. It feels one-directional in some circumstances, but I'm not an expert...
                            # The remaining relations are obviously symmetric:
                            # Association, Positive/Negative Correlation, Comparison, Co-Treament, and Drug Interaction.
                            # The only DocRED relation marked symmetric is "sister city".
                            # "spouse" and "sibling" should also be marked as such, though, so that might get updated.
                            # (Those relations aren't examined in these experiments)
                            if rel_info[prompt_id]["symmetric"] == "true":
                                ans.extend([(a[1], a[0]) for a in ans if (a[1], a[0]) not in ans])
                            # print(f"{ans=}")
                            # No sense in setting up a bunch of examples if none are correct.
                            # Maybe a retrieval system (RAG?) can make this selection in the wild?
                            if len(ans) == 0:
                                print(f"{prompt_id=} has noans")
                                continue
                            print(f"Document {p_doc.num} for {nonlin} {pooler} {prompt_id} at {nps} passes.", flush=True)

                            if big_batch:
                                from time import time
                                # stime = time()
                                # scores_old = rme_inner(fb, ans, nps, scorers, fuzzy_embeds, e_to_m_map, ev_tkns, prompt_data, prompt_id)
                                # otime += time() - stime
                                # stime = time()
                                scores = rme_inner2(fb, rel_info, ans, nps, scorers, fuzzy_embeds, e_to_m_map, ev_tkns, prompt_data, prompt_id, max_batch)
                                # ntime += time() - stime
                                # # print(f"{'Original' if dtime < dtime2 else 'New'}: {dtime:.2f} vs {dtime2:.2f}")
                                # with open(f'res/final/{task_name}_{model}_{dset}_{p_doc.num}_{num_blanks}b_{nps}p.pickle', 'rb') as other_scores_pickle:
                                #     scores_c = pickle.load(other_scores_pickle)
                                # # scores_a = sorted(scores_old['pll'], key=lambda x: (x[-1], *x[0:3]))
                                # # scores_b = sorted(scores['pll'], key=lambda x: (x[-1], *x[0:3]))
                                # # scores_c = sorted(scores_c[nonlin][pooler][prompt_id]['pll'], key=lambda x: (x[-1], *x[0:3]))

                                # scores_a = sorted(scores_old['pll'], key=lambda x: tuple(x[0:2]))
                                # scores_b = sorted(scores['pll'], key=lambda x: tuple(x[0:2]))
                                # scores_c = sorted(scores_c[nonlin][pooler][prompt_id]['pll'], key=lambda x: tuple(x[0:2]))

                                # # # print([(a, b, c, d) for a, b, c, d in scores_a])
                                # # # print([(a, b, c, d) for a, b, c, d in scores_b])
                                # def printem(t):
                                #     return f"({t[0]},{t[1]},{t[3]:.4f})"

                                # for i, (a, b, c) in enumerate(zip(scores_a, scores_b, scores_c)):
                                # #     # assert a[0] == b[0], f"{i}: {a} <-> {b}"
                                # #     # assert a[1] == b[1], f"{i}: {a} <-> {b}"
                                # #     # assert a[2] == b[2], f"{i}: {a} <-> {b}"
                                # #     # assert a[2] == c[2], f"{i}: {a} <-> {c}"
                                #     print(f"{i: 4}: {printem(a)}, {printem(b)}, {printem(c)}")
                                # print(f"{'Original' if otime < ntime else 'New'}: {otime:.2f} vs {ntime:.2f}")
                                # exit(0)

                                
                            else:
                                                    # fb: FitBert, ans, nps, nonlins, scorers, fuzzy_embeds, e_to_m_map, ev_tkns, prompt_data, prompt_id, max_batch
                                scores = rme_inner(fb, ans, nps, None, scorers, fuzzy_embeds, e_to_m_map, ev_tkns, prompt_data, prompt_id, max_batch)
                                # print([(a, b, c, d) for a, b, c, d in scores['pll']])
                                
                            
                            all_scores[nonlin][pooler][prompt_id] = scores
                            # return
                #exit(0)
                with open(docfile.replace(f'_{num_passes}p', f'_{nps}p'), 'wb') as resfile:
                    pickle.dump(all_scores, resfile)
                # if os.path.getsize(stopfile) > 0:
                #     break
            #     scores_a = sorted(all_scores['top10'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
            #     scores_c = sorted(all_scores['top20'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
            #     scores_d = sorted(all_scores['softmax'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
            #     # scores_f = sorted(all_scores['top20'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
            #     scores_f = sorted(all_scores[None][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
            #     scores_h = sorted(all_scores['top50'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
            #     # scores_i = sorted(all_scores['top20'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
            #     # with open('res/test/docred_bert-large-cased_dev_0_0b_0p.pickle', 'rb') as f:
            #     #     scores_two = pickle.load(f)
            #     with open('res/test/docred_bert-large-cased_dev_0_0b_0p.pickle', 'rb') as f:
            #         scores_old = pickle.load(f)
            #     scores_b = sorted(scores_old['top10'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
            #     scores_i = sorted(scores_old['top50'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
            #     scores_e = sorted(scores_old['softmax'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
            #     scores_g = sorted(scores_old[None][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))

            #     for z, (a, b, c, d, e, f, g, h, i) in enumerate(zip(scores_a, scores_b, scores_c, scores_d, scores_e, scores_f, scores_g, scores_h, scores_i)):
            #         if z % 10 == 0:
            #             print( "-------------------------------------------------------------------------------------")
            #             print(f"new@10 | old@10 | new@20 |X| new@SM | old@SM |X|  new@N |  OLD@N |X| new@50 | old@50 ")
            #         print(f"{a[3]:.3f} | {b[3]:.3f} | {c[3]:.3f} |X| {d[3]:.3f} | {e[3]:.3f} |X| {f[3]:.3f} | {g[3]:.3f} |X| {h[3]:.3f} | {i[3]:.3f}")
            # exit(0)

            # if p_doc.num == 9:
            #     print(f"{'Original' if otime < ntime else 'New'}: {otime:.2f} vs {ntime:.2f}")
            #     exit(0)
                    
        # if stopfile:
        #     with open('stopper.txt', 'w') as stopfile:
        #         pass


def run_many_experiments(p_doc, docfile, rel_info, nonlins, scorers, num_blanks, num_passes, max_batch=2000, big_batch=False):
    # import time
    with torch.no_grad():
        for nps in range(0, num_passes):
            scores = None
            print(f"{nps=}")
            if num_blanks > 0 and nps == 0:
                continue
            cur_docfile = docfile.replace(f'_{num_passes}p', f'_{nps}p')
            # cur_docfile = cur_docfile.replace('test', 'final')
            if os.path.exists(cur_docfile):
                continue
            else:
                if nps == 0:
                    nls = [None]  # The only valid nonlinearity for the first pass is no nonlinearity.
                else:
                    nls = nonlins
            evs, ev_tkns = p_doc.entity_vecs(nonlinearity=None, pooling=None, passes=nps)
            # Problem is here:
            fuzzy_embeds = {nl: {e:[output_to_fuzzy_embeddings(fb, fb.nonlin(nl)(v1).unsqueeze(0)) for v1 in evs[e]] for e in evs } for nl in nls}
            # exit(0)
            e_to_m_map = p_doc.masked_doc['ents']
            # print(f"A: {torch.cuda.mem_get_info()}")
            # times = []
            # Check to see if we've run this experiment.


            # print("RESETTING SCORES")
            for prompt_id in prompts: # rel_info:
                print(f"{prompt_id=}")
                if prompt_id.split('|')[0] not in p_doc.relations:
                    continue
                # print(p_doc.entities)
                # nnow = time.time()
                # print(f"PI: {prompt_id}")
                ans = [(a[0], a[2]) for a in p_doc.answers(detailed=False) if a[1] == prompt_id.split('|')[0]]
                # BioRED explicitly states that all relations are non-directional.
                # This is honestly false, but we marked the ones that aren't clearly non-directional to avoid issues.
                # For example, "Conversion" is a one-way process between chemicals.
                # "Bind" is questionable in this regard. It feels one-directional in some circumstances, but I'm not an expert...
                # The remaining relations are obviously symmetric:
                # Association, Positive/Negative Correlation, Comparison, Co-Treament, and Drug Interaction.
                # The only DocRED relation marked symmetric is "sister city".
                # "spouse" and "sibling" should also be marked as such, though, so that might get updated.
                # (Those relations aren't examined in these experiments)
                if rel_info[prompt_id]["symmetric"] == "true":
                    ans.extend([(a[1], a[0]) for a in ans if (a[1], a[0]) not in ans])
                # print(f"{ans=}")
                # No sense in setting up a bunch of examples if none are correct.
                # Maybe a retrieval system (RAG?) can make this selection in the wild?
                if len(ans) == 0:
                    print(f"{prompt_id=} has noans")
                    continue
                print(f"Document {p_doc.num} for {prompt_id} at {nps} passes.", flush=True)

                if big_batch:
                    scores = rme_inner2(fb, rel_info, scores, ans, nps, nls, scorers, fuzzy_embeds, e_to_m_map, ev_tkns, p_doc.entity_types, prompt_data, prompt_id, max_batch)
                else:
                    scores = rme_inner(fb, ans, nps, nls, scorers, fuzzy_embeds, e_to_m_map, ev_tkns, prompt_data, prompt_id, max_batch)

            with open(docfile.replace(f'_{num_passes}p', f'_{nps}p'), 'wb') as resfile:
                print(f"Saving to: {docfile.replace(f'_{num_passes}p', f'_{nps}p')}")
                pickle.dump(scores, resfile)


def rme_inner(fb: FitBert, ans, nps, nonlins, scorers, fuzzy_embeds, e_to_m_map, ev_tkns, prompt_data, prompt_id, max_batch) -> dict:
    scores = {}
    for sc in scorers:
        scores[sc.label] = []
    torch.cuda.empty_cache()
    # res = defaultdict(lambda:-float('inf'))
    all_replaced_m = {}
    all_labels_m = {}
    all_pins_m = {}
    for e1 in fuzzy_embeds:
        for e2 in fuzzy_embeds:
            if e1 != e2:
                _seen = set()
                for v1, m1 in zip(fuzzy_embeds[e1], e_to_m_map[e1]):
                    for v2, m2 in zip(fuzzy_embeds[e2], e_to_m_map[e2]):
                        if nps == 0:
                            vals = tuple(ev_tkns[m1].tolist() + [None] + ev_tkns[m2].tolist())
                            if vals in _seen:
                                continue
                            else:
                                _seen.add(vals)
                        for pin, pe in enumerate(prompt_data[prompt_id]):
                            rep_vecs = replace_embeddings(v1, v2, pe)
                            mv, ms = mask_vectors(fb, rep_vecs, keep_original=True, add_special_tokens=True)
                            size = mv.shape[0]
                            if size not in all_replaced_m:
                                all_replaced_m[size] = []
                                all_labels_m[size] = []
                                all_pins_m[size] = []
                            all_replaced_m[size].append(mv)
                            all_labels_m[size].append((e1, e2, m1, m2))
                            all_pins_m[size].append(pin)
    for size in all_replaced_m:
        print(f"{size=}")
        _max_batch_resized = max_batch - (max_batch % size)
        _sentences_per_batch = _max_batch_resized // size
        all_labels = all_labels_m[size]
        print(f"{len(all_labels)} candidate statements.", flush=True)
        all_pins = all_pins_m[size]
        bert_forward = torch.cat(all_replaced_m[size], dim=0).cpu()
        for v in all_replaced_m[size]:
            del v
        # all_replaced_m[size] = None
        # torch.cuda.empty_cache()
        # print(bert_forward.shape)
        # fwd_pieces = []
        print(f"Scoring bert forward ({size}, {len(bert_forward)})", flush=True)
        while len(bert_forward) > 0:
            # print(f"BF: {len(bert_forward)} ({min(_max_batch_resized, len(bert_forward))//size}/{len(all_labels)})", flush=True)
            # print(bert_forward[:_max_batch_resized].to(device=fb.device, dtype=fb.float_dtype)[0].shape)
            # print(bert_forward[:_max_batch_resized].to(device=fb.device, dtype=fb.float_dtype)[0])
            # print(fb.bert(inputs_embeds=bert_forward[:_max_batch_resized].to(device=fb.device, dtype=fb.float_dtype)).logits[0].shape)
            # print(fb.bert(inputs_embeds=bert_forward[:_max_batch_resized].to(device=fb.device, dtype=fb.float_dtype)).logits[0])
            # sm = fb.softmax_(fb.bert(inputs_embeds=bert_forward[:_max_batch_resized].to(device=fb.device, dtype=fb.float_dtype)).logits)[0]
            # print(sm.shape)
            # print(sm)
            # exit(0)
            with torch.autocast("cuda", dtype=fb.float_dtype):
                sm_bert = fb.softmax_(fb.bert(inputs_embeds=bert_forward[:_max_batch_resized].to(device=fb.device, dtype=fb.float_dtype)).logits.detach())[:, 1:, :]
            bert_forward = bert_forward[_max_batch_resized:]
            # torch.cuda.empty_cache()
            # sm_bert = fb.softmax_(fb.bert(inputs_embeds=bert_forward.to(fb.device)).logits)[:, 1:, :]
            # print(f"Document {p_doc.num} Past.", flush=True)
            # This will cause issues.
            # sm_bert = torch.cat(fwd_pieces) if len(fwd_pieces) > 1 else fwd_pieces[0]
            # print(len(sm_bert.view(-1, score_len, sm_bert.shape[1], sm_bert.shape[2])), len(all_labels))
            # print(all_labels)
            # print(f"SM: {sm_bert.shape}")
            # print(len(sm_bert.view(-1, score_len, sm_bert.shape[1], sm_bert.shape[2])), len(all_labels))
            for (e1, e2, m1, m2), pin, s in zip(all_labels[:_sentences_per_batch], all_pins[:_sentences_per_batch], sm_bert.view(-1, size, sm_bert.shape[1], sm_bert.shape[2])):
                # print("S:", s.shape)
                for scorer in scorers:
                    origids = None
                    if scorer.label == "pll":
                        # Then make the index from its parts, same as the other thing.
                        origids = replace_ids(ev_tkns[m1], ev_tkns[m2], prompt_data[prompt_id][pin])[1:-1].to(fb.device)
                    scores[scorer.label].append((e1, e2, (e1, e2) in ans, scorer(s, origids=origids)))
            all_labels = all_labels[_sentences_per_batch:]
            all_pins = all_pins[_sentences_per_batch:]
            # print(f"Document {p_doc.num} Tick.", flush=True)
            # for scorer in scorers:
            #     print(scorer.label, len(all_scores[nonlin][pooler][prompt_id][scorer.label]))
            del s
            del sm_bert
            # torch.cuda.empty_cache()
        del bert_forward
    return scores

def rme_inner2(fb: FitBert, rel_info, scores:dict, ans, nps, nonlins, scorers, fuzzy_embeds, e_to_m_map, ev_tkns, entity_types, prompt_data, prompt_id, max_batch) -> dict:
    if not scores:
        scores = {}
    for nl in nonlins:
        if nl not in scores:
            scores[nl] = {None:{prompt_id: {s.label:[] for s in scorers}}}
        else:
            if prompt_id not in scores[nl][None]:
                scores[nl][None][prompt_id] = {s.label:[] for s in scorers}
            else:
                for s in scorers:
                    if s.label not in scores[nl][None][prompt_id]:
                        scores[nl][None][prompt_id][s.label] = []
    # scores = {nl: {None:{prompt_id: {s.label:[] for s in scorers}}} for nl in nonlins}
    # for sc in scorers:
    #     scores[sc.label] = []
    # torch.cuda.empty_cache()
    # res = defaultdict(lambda:-float('inf'))
    all_replaced_m = {}
    all_masks_m = {}
    all_labels_m = {}
    all_sizes_m = {}
    all_nls_m = {}
    all_pins_m = {}
    maxlen = len(max(ev_tkns.values(), key=len)) * 2 + max(len(p['input_ids']) for p in prompt_data[prompt_id]) - 2
    # print(ev_tkns.values())
    # print(prompt_data[prompt_id])
    # print(list(p['input_ids'] for p in prompt_data[prompt_id]))
    # print(list(fb.tokenizer.convert_ids_to_tokens(p['input_ids']) for p in prompt_data[prompt_id]))
    # print(len(max(ev_tkns.values(), key=len)))
    # print(max(len(p['input_ids']) for p in prompt_data[prompt_id]))
    # print(maxlen)
    # exit(0)

    # print(f"start: {maxlen=}")
    for nl in nonlins:
        for e1 in fuzzy_embeds[nl]:
            if not any(t in rel_info[prompt_id]['domain'] for t in entity_types[e1]):
                continue
            # print("OK")
            for e2 in fuzzy_embeds[nl]:
                if not any(t in rel_info[prompt_id]['range'] for t in entity_types[e2]):
                    continue
                if e1 != e2:
                    _seen = set()
                    for v1, m1 in zip(fuzzy_embeds[nl][e1], e_to_m_map[e1]):
                        for v2, m2 in zip(fuzzy_embeds[nl][e2], e_to_m_map[e2]):
                            if nps == 0:
                                vals = tuple(ev_tkns[m1].tolist() + [None] + ev_tkns[m2].tolist())
                                if vals in _seen:
                                    continue
                                else:
                                    _seen.add(vals)
                            for pin, pe in enumerate(prompt_data[prompt_id]):
                                rep_vecs = replace_embeddings(v1, v2, pe)
                                mv, ms = mask_vectors(fb, rep_vecs, keep_original=True, add_special_tokens=True, pad_to=0)
                                # print(mv.shape, ms.shape)

                                size = mv.shape[1]
                                # print(fb.bert(inputs_embeds=mv.to(fb.device), attention_mask=ms.to(fb.device)).logits[:, 1:, :])
                                if size not in all_replaced_m:
                                    all_replaced_m[size] = []
                                    all_masks_m[size] = []
                                    all_labels_m[size] = []
                                    all_sizes_m[size] = []
                                    all_nls_m[size] = []
                                    all_pins_m[size] = []
                                all_replaced_m[size].append(mv)
                                all_masks_m[size].append(ms)
                                all_sizes_m[size].append(len(mv))
                                all_labels_m[size].append((e1, e2, m1, m2))
                                all_nls_m[size].append(nl)
                                all_pins_m[size].append(pin)
    # print("end")
    for size in all_replaced_m:
        print(f"{size=}")
        # _max_batch_resized = max_batch - (max_batch % size)
        # TODO: cumsum this to find out how many sentences per batch we actually use.
        # cumsum
        # where > max_batch
        # slice to where - 2
        # _sentences_per_batch = _max_batch_resized // size

        all_labels = all_labels_m[size]
        all_sizes = all_sizes_m[size]
        all_nls = all_nls_m[size]
        all_pins = all_pins_m[size]
        print(f"{len(all_labels)} candidate statements.", flush=True)

        # print([v.shape for v in all_replaced_m[size]])

        bert_forward = torch.cat(all_replaced_m[size], dim=0).cpu()
        for v in all_replaced_m[size]:
            del v
        bert_forward_mask = torch.cat(all_masks_m[size], dim=0).cpu()
        # print(sum(all_sizes))
        # print(bert_forward.shape)
        # print(bert_forward_mask.shape)
        for v in all_masks_m[size]:
            del v
        # all_replaced_m[size] = None
        torch.cuda.empty_cache()
        # print(bert_forward.shape)
        # fwd_pieces = []
        print(f"Scoring bert forward ({size}, {len(bert_forward)})", flush=True)
        stride_length = 65536 // size
        if "roberta" in fb.model_name.lower() or "modern" in fb.model_name.lower():  # Larger models support smaller batches
            stride_length /= 2
        if "biobert" in fb.model_name.lower():
            stride_length /= 2
        # if fb.modern:  # ModernBERT can handle big batches but is very, very slow. Maybe smaller batches work?
        #     stride_length /= 2

        strides = batchify(all_sizes, stride_length)
        # print(strides)
        # tst = []
        # for st in strides:
        #     tst.append(sum(all_sizes[:st]))
        #     all_sizes = all_sizes[st:]
        # print(tst)
        # print(sum(tst))
        # exit(0)
        # print()
        while len(bert_forward) > 0:

            stride, *strides = strides
            sizes = all_sizes[:stride]
            nls = all_nls[:stride]
            pins = all_pins[:stride]
            all_sizes = all_sizes[stride:]
            all_nls = all_nls[stride:]
            all_pins = all_pins[stride:]
            this_batch = sum(sizes)
            lg = max(sizes) + 1

            print(f"BF: {len(all_labels)} statements left, ({this_batch}/{len(bert_forward)} sequences)", flush=True)
            with torch.autocast("cuda", dtype=fb.float_dtype):
                sm_bert = fb.softmax_(fb.bert(inputs_embeds=bert_forward[:this_batch, :lg].to(fb.device), attention_mask=bert_forward_mask[:this_batch, :lg].to(fb.device)).logits)[:, 1:, :]
            bert_forward = bert_forward[this_batch:]
            bert_forward_mask = bert_forward_mask[this_batch:]
            # torch.cuda.empty_cache()
            labels = all_labels[:stride]

            for (e1, e2, m1, m2), size, nl, pin in zip(labels, sizes, nls, pins):
                # s = sm_bert[:size]
                # print("S:", s.shape)
                for scorer in scorers:
                    origids = None
                    if scorer.label == "pll":
                        # Then make the index from its parts, same as the other thing.
                        origids = replace_ids(ev_tkns[m1], ev_tkns[m2], prompt_data[prompt_id][pin])[1:-1].to(fb.device)
                    scores[nl][None][prompt_id][scorer.label].append((e1, e2, (e1, e2) in ans, scorer(sm_bert[:size], origids=origids)))
                sm_bert = sm_bert[size:]
            # print(sm_bert.shape)
            all_labels = all_labels[stride:]
            # all_sizes = all_sizes[_sentences_per_batch:]
            # print(f"Document {p_doc.num} Tick.", flush=True)
            # for scorer in scorers:
            #     print(scorer.label, len(all_scores[nonlin][pooler][prompt_id][scorer.label]))
            # del s
            # print("Del sm_bert")
            del sm_bert
            # print("Empty cache")
            # torch.cuda.empty_cache()
        # print("Del bert_forward")
        del bert_forward
    # print("Return.")
    return scores
    
def batchify(values, max_per_batch):
    tot = 0
    subsets = []
    num = 0
    for v in values:
        if v + tot > max_per_batch:
            subsets.append(num)
            tot = v
            num = 1
        else:
            tot += v
            num += 1
    subsets.append(num)
    return subsets


def dashsplit(nums):
    parts = nums.split("-")
    assert len(parts) <= 2
    if len(parts) == 1:
        num_blanks_start, num_blanks = 0, int(nums or 0)
    else:
        num_blanks_start = int(parts[0] or 0)
        num_blanks = int(parts[1] or 3)
    return num_blanks_start, num_blanks


def report(message: str, last_message: int) -> int:
    try:
        if last_message >= 0:
            requests.post(f"https://api.telegram.org/bot5651481224:AAGWfTQsv2xe29mJNNsG2ocGvhKQIzA8n30/deleteMessage?chat_id=-1003527304165&message_id={last_message}")
            time.sleep(.5)
    except:
        # Assume that the last message failed to delete. However, also assume that there's a temporary problem and don't try to post a new message.
        # Next time, try to delete the prior message again.
        return last_message
    try:
        response = requests.post(f"https://api.telegram.org/bot5651481224:AAGWfTQsv2xe29mJNNsG2ocGvhKQIzA8n30/sendMessage?chat_id=-1003527304165&text='{message}'")
        rj = response.json()
        if "result" in rj and "message_id" in rj["result"]:
            return int(rj["result"]["message_id"])
        else:
            return -1
    except:
        return -1

    # dset = sys.argv[1]
    # vals = [1, 5, 2, 6, 3, 8, 5, 2, 5, 3, 6, 8, 3, 1, 2, 7 , 6, 2, 4, 5, 4, 7]
    # bats = batchify(vals, 10)
    # print(vals)
    # print(bats)
    # for bat in bats:
    #     vs = vals[:bat]
    #     vals = vals[bat:]
    #     print(sum(vs), vs)
    # exit(0)


if __name__ == '__main__':
    import sys
    num_blanks = sys.argv[1] if len(sys.argv) > 1 else "2"
    num_passes = int(sys.argv[2]) if len(sys.argv) > 2 else 3
    data_path = sys.argv[3] if len(sys.argv) > 3 else "data"
    task_name = sys.argv[4] if len(sys.argv) > 4 else "docred"
    data_set = sys.argv[5] if len(sys.argv) > 5 else "dev"
    resdir = sys.argv[6] if len(sys.argv) > 6 else "res/testtest"
    max_batch = int(sys.argv[7]) if len(sys.argv) > 7 else 2000
    use_ent = (int(sys.argv[8]) == 1) if len(sys.argv) > 8 else False
    stopfile = sys.argv[9] if len(sys.argv) > 9 else "stopper.txt"
    model = sys.argv[10] if len(sys.argv) > 10 else "modernbert"
    start_at = sys.argv[11] if len(sys.argv) > 11 else "0"
    import socket
    c_name = sys.argv[12] if len(sys.argv) > 12 else socket.gethostname()

    num_blanks_start, num_blanks = dashsplit(num_blanks)
    start_at, end_at = dashsplit(start_at)

    # SENTS_PER_BATCH = int(sys.argv[5]) if len(sys.argv) > 5 else 256
    # BATCH_SIZE = 512
    # # If you have memory issues with some documents, you can add them here to skip. 332 is particularly difficult.
    # TOO_BIG = []  # [332]
    # SENTS_PER_BATCH = 32
    # width = 2
    # resdir = 'res'
    # doc=2
    # num_blanks=2
    # skip = [73, 234]
    # num_passes=3
    # nonlins = ["relu", "softmax", None]
    # poolers = ["mean", "max", None]
    # scorers = [CSD(), ESD(), JSD(), MSD(), HSD()]

    models = {
        "bert": "bert-large-cased",
        "bert-base": "bert-base-cased",
        "bert-uncased": "bert-large-uncased",
        "bert-base-uncased": "bert-base-uncased",
        "roberta-base": "roberta-base",  # (cased)
        "roberta": "roberta-large",  # (cased)
        "distilbert": "distilbert-base-cased",  # Doesn't work for some reason? Ranking looks suspicious and follows no trends.
        "biobert": "dmis-lab/biobert-large-cased-v1.1",
        "biobert-base": "dmis-lab/biobert-base-cased-v1.1",
        # "biobert-base": "dmis-lab/biobert-base-uncased-v1.1",
        "pubmedbert": "microsoft/BiomedNLP-PubMedBERT-large-uncased-abstract",
        "medbert": "Charangan/MedBERT",  # It should work! We just didn't get around to testing it in time...
        "clinical_biobert": "emilyalsentzer/Bio_ClinicalBERT",
        # "modernbert": "answerdotai/ModernBERT-large",
        # "modernbert-large": "answerdotai/ModernBERT-large",
        # "modernbert-base": "answerdotai/ModernBERT-base",
    }

    for model in models:
        s = "Hello, world!"
        tk = AutoTokenizer.from_pretrained(models[model])
        print(model)
        print("  ", tk.convert_ids_to_tokens(tk(s)['input_ids']))
        m = AutoModelForMaskedLM.from_pretrained(models[model])
        print(f"  {sum(p.numel() for p in m.parameters() if p.requires_grad)/1000000:.0f}M")
    exit(0)


    if model.lower() in models:
        model = models[model.lower()]

    with open(f'{data_path}/{task_name}/rel_info_full.json', 'r') as rel_info_file:
        rel_info = json.load(rel_info_file)

    assert task_name in os.listdir('./data'), f"Data set missing for task '{task_name}'"
    
    skip = []
    if task_name == "biored":
        if num_blanks > 0:
            # skip += [38, 78, 87, 96, 136, 150, 151, 163]
            pass
        if num_blanks > 1:
            pass
            # skip += [20, 32]
    elif "dwie" in task_name:
        skip += [39, 170, 428, 699]  # DW_17797401
    else:
        skip=[723, 777]
        if num_blanks >= 4:
            skip += [121]

    # def run_many_experiments(data, dset, rel_info, nonlins, poolers, scorers, resdir, num_blanks, num_passes=1, max_batch=2000, use_ent=True, skip=[], stopfile='', model='bert-large-cased'):

    os.makedirs(resdir, exist_ok=True)

    torch.set_printoptions(edgeitems=4)

    reduced_precision = False
    if "modern" in model.lower():
        reduced_precision = True

    if reduced_precision:
        torch.set_float32_matmul_precision("high")

    # try:
    #     os.remove('res/test/docred_bert-large-cased_dev_0_0b_0p.pickle', )
    # except:
    #     pass

    import sys
    print(sys.version)
    # print(sys.version_info)

    # for x, y in FitBert.nonlins.items():
    #     print(x, callable(y), type(y))
    # exit(0)

    # _nps = 2
    # if os.path.exists(f'res/test3/docred_bert-large-cased_dev_0_0b_{_nps}p.pickle'):
    #     print("Loading from files.")
    #     with open(f'res/test3/docred_bert-large-cased_dev_0_0b_{_nps}p.pickle', 'rb') as f:
    #         scores = pickle.load(f)

    #     # print(scores.keys())
    #     # print(scores[None][None].keys())
    #     scores_a = sorted(scores['top10'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
    #     scores_c = sorted(scores['top20'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
    #     scores_d = sorted(scores['softmax'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
    #     # scores_f = sorted(all_scores['top20'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
    #     scores_f = sorted(scores[None][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
    #     scores_h = sorted(scores['top50'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
    #     # scores_i = sorted(all_scores['top20'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
    #     # with open('res/test/docred_bert-large-cased_dev_0_0b_0p.pickle', 'rb') as f:
    #     #     scores_two = pickle.load(f)
    #     with open(f'res/test2/docred_bert-large-cased_dev_0_0b_{_nps}p.pickle', 'rb') as f:
    #         scores_old = pickle.load(f)
    #     scores_b = sorted(scores_old['top10'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
    #     scores_i = sorted(scores_old['top50'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
    #     scores_e = sorted(scores_old['softmax'][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))
    #     scores_g = sorted(scores_old[None][None]['P17']['pll'], key=lambda x: tuple(x[0:2]))

    #     for z, (a, b, c, d, e, f, g, h, i) in enumerate(zip(scores_a, scores_b, scores_c, scores_d, scores_e, scores_f, scores_g, scores_h, scores_i)):
    #         if z % 10 == 0:
    #             print( "-------------------------------------------------------------------------------------")
    #             print(f"new@10 | old@10 | new@20 |X| new@SM | old@SM |X|  new@N |  OLD@N |X| new@50 | old@50 ")
    #         print(f"{a[3]:.3f} | {b[3]:.3f} | {c[3]:.3f} |X| {d[3]:.3f} | {e[3]:.3f} |X| {f[3]:.3f} | {g[3]:.3f} |X| {h[3]:.3f} | {i[3]:.3f}")
    #     exit(0)

    # fb = extend_bert(FitBert(model_name="bert-large-cased", reduced_precision=False), 50, 0, roberta=False)

    with torch.no_grad():
        torch.cuda.empty_cache()
        fb = extend_bert(FitBert(model_name=model, reduced_precision=reduced_precision), 200, num_blanks)
        # print(fb.tokenizer.convert_ids_to_tokens(fb.entity_tokens))
        model = model.split('/')[-1]
        # nls = {None: lambda x:x, "softmax": FitBert.softmax, "relu":torch.relu}
        # global one_hot_punctuation
        # one_hot_punctuation = read_punctuation(fb)

        # is_re = task_name.startswith("re-")
        if task_name[:3] == "re-":
            task_type = task_name[3:]
        else:
            task_type = task_name
        
        if "docred" in task_type or "docshred" in task_type:
            prompts = ['P17', 'P27', 'P131', 'P150', 'P161', 'P175', 'P527', 'P569', 'P570', 'P577']
        elif task_type == "biored":
            prompts = ['Association', 'Bind', 'Negative_Correlation', 'Positive_Correlation']
        elif task_type == "dwie":
            prompts = ['gpe', 'in', 'citizen_of', 'based_in', 'member_of',
                        'agent_of', 'head_of', 'head_of_state', 'agency_of',
                          'head_of_gov']  # ['minister_of', 'event_in', 'part_of', 'institution_of', 'appears_in', 'award_received', 'ministry_of', 'player_of', 'vs', 'created_by']
            for p in prompts:
                assert p in rel_info, f"{p} not found, might be spelled wrong."
        else:
            prompts = list(sorted(rel_info.keys()))
        prompt_data = get_prompt_info(prompts, rel_info, fb)

    from itertools import count
    import time

    def work():
        last_message = -1
        tnow = int(time.time())
        try:
            for doc in count(start=start_at):
                if doc > end_at:
                    break
                did_work = False
                for _nb in range(num_blanks_start, num_blanks + 1):
                    try:
                        p_doc, docfile = process_single_document(data_path, resdir, fb=fb, task_name=task_name, dset=data_set, doc=doc,
                                                                num_blanks=_nb, num_passes=num_passes, use_ent=use_ent, top_k=0)
                        if not p_doc or not docfile:
                            continue
                        did_work = True
                        #(p_doc, docfile, rel_info, nonlins, scorers, num_blanks, num_passes, max_batch=2000, big_batch=False):
                        run_many_experiments(
                            p_doc=p_doc,
                            docfile=docfile,
                            rel_info=rel_info,
                            nonlins=["top10"],#[None, "softmax", "top5", "top10", "top25", "top50", "top100"],
                            #poolers=[None],
                            scorers=[PLL()],
                            # resdir=resdir,
                            num_blanks=_nb,
                            num_passes=num_passes,
                            max_batch=max_batch,
                            # use_ent=use_ent,
                            # skip=skip,
                            # stopfile=stopfile,
                            # model=model,
                            # start_at=start_at,
                            # reduced_precision=reduced_precision,
                            big_batch=True
                        )
                    except StopIteration:

                        print("Done!")
                        report(f'{c_name} has finished all docs!:\n`{task_name}({data_set}) {model} {resdir}`', last_message)
                        return
                    except KeyboardInterrupt:
                        print("Stopping!")
                        report(f'{c_name} has been paused at doc {doc}!:\n`{task_name}({data_set}) {model} {resdir}`', last_message)
                        return
                if did_work:

                    tdiff = int(time.time()) - tnow
                    # if doc%10 == 0:
                    last_message = report(f'{c_name} has finished doc {doc}. It took {tdiff} seconds.', last_message)
                    tnow = int(time.time())
        except BaseException as e:
            report(f'{c_name} has FAILED.\n`{task_name}({data_set}) {model} {resdir}`\n{Exception}: {e}', last_message)
            print(e)
            raise e

    # def work():
    #     for doc in count(start=start_at):
    #         for _nb in [1, 2]:
    #             p_docA, _ = process_single_document(data_path, resdir, fb=fb, task_name=task_name, dset=data_set, doc=doc,
    #                                                     num_blanks=_nb, num_passes=2, use_ent=use_ent, top_k=0)
    #             # print(task_name)
    #             if not p_docA:
    #                 continue
    #             p_docB, _ = process_single_document(data_path, resdir, fb=fb, task_name=task_name.replace("4421", "4423"), dset=data_set, doc=doc,
    #                                                     num_blanks=_nb, num_passes=2, use_ent=use_ent, top_k=0)
    #             # print(task_name.replace("4421", "4423"))
    #             if len(p_docA._mask_tkns) != len(p_docB._mask_tkns):
    #                 print(f"Doc {doc} failed mask length check @ {_nb=}.")
    #                 print(p_docA._ent_tkns)
    #                 print(p_docB._ent_tkns)
    #                 print(p_docA._mask_tkns)
    #                 print(p_docB._mask_tkns)
    #             if any(a != b for a, b in zip(p_docA._mask_tkns, p_docB._mask_tkns)):
    #                 print(f"Doc {doc} failed mask identity check @ {_nb=}.")
    #                 print(p_docA._ent_tkns)
    #                 print(p_docB._ent_tkns)
    #                 print(p_docA._mask_tkns)
    #                 print(p_docB._mask_tkns)
    work()
    # Note to self: Why in the heck are top50 and top100 identical, but softmax and top10 not?!
    # Answer: Because some of the methods were not being seen as callable methods, and were defaulting to the identity function.