Upload ContextualDocumentEmbeddingTransformer

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

config.json

@@ -0,0 +1,28 @@
+{
+  "_name_or_path": "/fsx-checkpoints/jxm/cde/2024-08-06-transductive-pretrain-transductive-long-10node-3/checkpoint-7176",
+  "architecture": "transductive",
+  "architectures": [
+    "ContextualDocumentEmbeddingTransformer"
+  ],
+  "attn_implementation": null,
+  "auto_map": {
+    "AutoConfig": "misc.ContextualModelConfig",
+    "AutoModel": "model.ContextualDocumentEmbeddingTransformer"
+  },
+  "cache_dir": null,
+  "config_name": null,
+  "disable_dropout": true,
+  "disable_transductive_rotary_embedding": true,
+  "embedder": "nomic-ai/nomic-bert-2048",
+  "embedder_rerank": "sentence-transformers/gtr-t5-base",
+  "embedding_output_dim": null,
+  "limit_layers": null,
+  "logit_scale": 50.0,
+  "max_seq_length": 512,
+  "model_revision": "main",
+  "tokenizer_name": null,
+  "torch_dtype": "float32",
+  "transductive_corpus_size": 512,
+  "transductive_sequence_dropout_prob": 0.0,
+  "transformers_version": "4.48.0.dev0"
+}

misc.py

@@ -0,0 +1,518 @@
+from typing import Dict, Iterable, List, Optional, Tuple, Union
+
+import collections
+import glob
+import json
+import hashlib
+import itertools
+import logging
+import multiprocessing
+import os
+import pickle
+import random
+import requests
+import sys
+import zipfile
+
+import datasets
+import numpy as np
+import torch
+import tqdm
+import transformers
+
+from cde.lib.dist import get_num_proc, get_rank
+
+
+def get_cde_cache_dir() -> str:
+    script_directory = os.path.normpath(
+        os.path.join(
+            os.path.dirname(os.path.abspath(__file__)),
+            os.pardir, os.pardir,
+        )
+    )
+    return os.path.join(script_directory, "data")
+
+
+def get_cache_location_from_kwargs(**kwargs):
+    cache_location = os.path.join(
+        get_cde_cache_dir(), "cluster"
+    )
+    os.makedirs(cache_location, exist_ok=True)
+    return os.path.join(cache_location, md5_hash_kwargs(**kwargs))
+
+
+def process_qrels_uncached(corpus: datasets.Dataset, qrels: datasets.Dataset) -> Tuple[Dict[str, List[float]], Dict[str, List[str]]]:
+    qrels_idxs = collections.defaultdict(list)
+    qrels_scores = collections.defaultdict(list)
+    corpus_ids = np.array(corpus['_id'])
+    skipped_qrels = 0
+
+    for ex in tqdm.tqdm(qrels, desc='processing qrels', colour='#964B00', leave=False):
+        #
+        # example:
+        # {
+        #  'query-id': 1, 
+        #  'corpus-id': 'b0680508-2019-04-18T13:48:51Z-00002-000',
+        #  'score': 2
+        # }
+        # 
+        q_id = str(ex['query-id'])
+        c_idxs = (corpus_ids == str(ex['corpus-id'])).nonzero()[0]
+        # 
+        assert len(c_idxs) <= 1, f"error - duplicate corpus ID? (found {len(c_idxs)} matches)"
+        # 
+        if len(c_idxs):
+            qrels_idxs[q_id].append(c_idxs[0])
+            qrels_scores[q_id].append(ex['score'])
+        else:
+            skipped_qrels += 1
+        #
+    
+    if skipped_qrels > 0:
+        logging.warning(f'Warning: Skipped {skipped_qrels}/{len(qrels)} qrels.')
+    
+    return qrels_idxs, qrels_scores
+
+
+def process_qrels(
+        corpus: datasets.Dataset, qrels: datasets.Dataset, 
+        use_cache: bool = True
+    ) -> Tuple[Dict[str, List[float]], Dict[str, List[str]]]:
+    dataset_cache_file = '_'.join(
+        (corpus.cache_files[0]['filename'], qrels.cache_files[0]['filename'])
+    )
+    cache_file = strip_extension(dataset_cache_file) + '_processed_qrels.p'
+    os.makedirs(os.path.dirname(cache_file), exist_ok=True)
+
+    if not (use_cache and os.path.exists(cache_file)):
+        qrels_idxs, qrels_scores = process_qrels_uncached(
+            corpus=corpus, qrels=qrels
+        )
+        if use_cache:
+            pickle.dump((qrels_idxs, qrels_scores), open(cache_file, 'wb'))
+    else:
+        qrels_idxs, qrels_scores = pickle.load(open(cache_file, 'rb'))
+    
+    return qrels_idxs, qrels_scores
+
+
+def strip_extension(filename: str) -> str:
+    """Strips file extension.
+
+    Ex:
+        >> strip_extension('/root/dir/sub/file.ext')
+        '/root/dir/sub/file'
+    """
+    return os.path.splitext(filename)[0]
+
+
+def md5_hash(t: Tuple[str]) -> str:
+    return hashlib.md5('__'.join(t).encode()).hexdigest()
+
+
+def md5_hash_kwargs(**kwargs) -> str:
+    # We ignore special hf args that start with _ like '__cached__setup_devices'.
+    safe_kwargs = {k: str(v) for k,v in kwargs.items() if not k.startswith('_')}
+    s = json.dumps(safe_kwargs, sort_keys=True)
+    return hashlib.md5(s.encode()).hexdigest()
+
+def download_url(url: str, save_path: str, chunk_size: int = 1024):
+    """Download url with progress bar using tqdm
+    https://stackoverflow.com/questions/15644964/python-progress-bar-and-downloads
+    Args:
+        url (str): downloadable url
+        save_path (str): local path to save the downloaded file
+        chunk_size (int, optional): chunking of files. Defaults to 1024.
+    """
+    r = requests.get(url, stream=True)
+    total = int(r.headers.get('Content-Length', 0))
+    with open(save_path, 'wb') as fd, tqdm.tqdm(
+        desc=save_path,
+        total=total,
+        unit='iB',
+        unit_scale=True,    
+        unit_divisor=chunk_size,
+    ) as bar:
+        for data in r.iter_content(chunk_size=chunk_size):
+            size = fd.write(data)
+            bar.update(size)
+
+
+def unzip(zip_file: str, out_dir: str):
+    print("unzipping =>", zip_file)
+    zip_ = zipfile.ZipFile(zip_file, "r")
+    zip_.extractall(path=out_dir)
+    zip_.close()
+
+
+def download_url_and_unzip(url: str, out_dir: str, chunk_size: int = 1024) -> str:
+    os.makedirs(out_dir, exist_ok=True)
+    dataset = url.split("/")[-1]
+    zip_file = os.path.join(out_dir, dataset)
+    
+    if not os.path.isfile(zip_file):
+        logging.info("Downloading {} ...".format(dataset))
+        download_url(url, zip_file, chunk_size)
+    
+    if not os.path.isdir(zip_file.replace(".zip", "")):
+        logging.info("Unzipping {} ...".format(dataset))
+        unzip(zip_file, out_dir)
+    
+    return os.path.join(out_dir, dataset.replace(".zip", ""))
+
+
+def tqdm_if_main_worker(iterable: Iterable, **kwargs) -> Iterable:
+    if get_rank() == 0:
+        return tqdm.tqdm(iterable, **kwargs)
+    else:
+        return iterable
+
+
+class ContextualModelConfig(transformers.configuration_utils.PretrainedConfig):
+    """We create a dummy configuration class that will just set properties
+    based on whatever kwargs we pass in.
+
+    When this class is initialized (see experiments.py) we pass in the
+    union of all data, model, and training args, all of which should
+    get saved to the config json.
+    """
+
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            try:
+                json.dumps(value)
+                setattr(self, key, value)
+            except TypeError:
+                # value was not JSON-serializable, skip
+                continue
+        super().__init__()
+
+
+def independent_crop(
+    input_ids: torch.Tensor, pad_token_id: int,
+    l1: int = 256, l2: int = 256) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Returns two independent crops from input_ids.
+    
+    Assumes input_ids has a beginning and end token, like 
+        [101, ..., 102, 0, 0, 0].
+
+    Args:
+        input_ids: tensor of IDs
+        pad_token_id: ID of pad tokens in input_ids
+        l1: length of span 1, cropped
+        l2: length of span 2, cropped
+    Returns:
+        span1: first crop (of length l1)
+        span2: second crop (of length l2)
+    """ 
+    # Count tokens until pad.
+    if (input_ids == pad_token_id).sum() == 0:
+        N = len(input_ids)
+    else:
+        N = (input_ids == pad_token_id).int().argmax().item()
+    
+    ####
+    ###
+    ##
+    ## Contriever:  We use the random cropping data
+    ## augmentation, with documents of 256 tokens and span 
+    ## sizes sampled between 5% and 50% of the document
+    ## length
+    ##
+    ###
+    #####
+    ####### LaPraDor: The maximum lengths set for queries and
+    ####### documents are 64 and 350...
+    #####
+    # TODO is this divide-by-two a good idea? (Don't want s1=s2 ever..)
+    nl1 = min(N//2, l1)
+    nl2 = min(N//2, l2)
+
+    s1_start = random.randint(1, N-nl1)
+    s2_start = random.randint(1, N-nl2)
+
+    s1_idxs = itertools.chain(
+        [0], range(s1_start, s1_start+nl1), [N-1]
+    )
+    s1 = input_ids[torch.tensor(list(s1_idxs))]
+    s2_idxs = itertools.chain(
+        [0], range(s2_start, s2_start+nl2), [N-1]
+    )
+    s2 = input_ids[torch.tensor(list(s2_idxs))]
+    return (s1, s2)
+
+
+def load_dataset_tables(
+    files: Iterable[str], num_workers: int = 16
+) -> Iterable[datasets.table.MemoryMappedTable]:
+    import concurrent
+    from multiprocessing import Pool
+
+    # num_workers = min(num_workers, len(files))
+    num_workers = min(32, len(files))
+
+    use_threads = True
+    if use_threads:
+        pool_cls = concurrent.futures.ThreadPoolExecutor
+        pool_kwargs = {"max_workers": num_workers}
+    else:
+        pool_cls = Pool
+        pool_kwargs = {"processes": num_workers}
+    
+    with pool_cls(**pool_kwargs) as pool:
+        if len(files) > 10:
+            files = tqdm_if_main_worker(
+                files,
+                desc=f"Loading {len(files)} files with {num_workers} workers",
+                total=len(files),
+                colour="#ffbd88"
+            )
+        
+        result = list(
+            pool.map(datasets.table.MemoryMappedTable.from_file, files)
+        )
+    return result
+
+
+def datasets_fast_load_from_disk(cache_path: str) -> datasets.Dataset:
+    logging.info(f"fast_load_from_disk called with path:", cache_path)
+    dataset_info_path = os.path.join(cache_path, "dataset_info.json")
+    with open(dataset_info_path, encoding="utf-8") as dataset_info_file:
+        dataset_info = datasets.DatasetInfo.from_dict(json.load(dataset_info_file))
+
+    dataset_state_path = os.path.join(cache_path, "state.json")
+    with open(dataset_state_path, encoding="utf-8") as state_file:
+        state = json.load(state_file)
+
+    files = glob.glob(os.path.join(cache_path, "data-*.arrow"))
+    files = sorted(files)
+    num_workers = get_num_proc()
+    ds_tables = load_dataset_tables(
+        files=files,
+        num_workers=num_workers
+    )
+    arrow_table = datasets.table.concat_tables(ds_tables)
+
+    split = state["_split"]
+    split = datasets.splits.Split(split) if split is not None else split
+
+    # print("returning dataset")
+    return datasets.Dataset(
+        arrow_table=arrow_table,
+        info=dataset_info,
+        split=split,
+        fingerprint=state["_fingerprint"],
+    )
+
+
+def tokenize_dataset(
+        dataset: datasets.Dataset,
+        tokenizer: transformers.PreTrainedTokenizer,
+        max_length: int,
+        text_key: str,
+        padding_strategy: str
+    ) -> datasets.Dataset:
+    def tokenize_text(ex: Dict) -> Dict:
+        tt = tokenizer(
+            ex[text_key],
+            max_length=max_length,
+            truncation=True,
+            padding=padding_strategy,
+        )
+        for k,v in tt.items():
+            ex[f"{text_key}_{k}"] = v
+        ex["length"] = [len(tt) for tt in ex[f"{text_key}_input_ids"]]
+        return ex
+
+    # generate unique hash for tokenizer
+    vocab = tokenizer.vocab
+    vocab_words = tuple(sorted(vocab.keys(), key=lambda word: vocab[word]))
+    vocab_hash = md5_hash(vocab_words)
+
+    data_fingerprint = '__'.join((
+        dataset._fingerprint, str(vocab_hash), str(max_length),
+        text_key, padding_strategy
+    ))
+    data_fingerprint = md5_hash(data_fingerprint)
+    dataset = dataset.map(
+        tokenize_text,
+        new_fingerprint=data_fingerprint,
+        batched=True,
+        load_from_cache_file=True,
+    )
+    return dataset
+
+
+class TensorRunningAverages:
+    _store_sum: Dict[str, torch.Tensor]
+    _store_total: Dict[str, torch.Tensor]
+
+    def __init__(self):
+        self._store_sum = {}
+        self._store_total = {}
+    
+    def __iter__(self) -> Iterable[str]:
+        return iter(self._store_sum.keys())
+
+    def update(self, key: str, val: Union[int, float, torch.Tensor]) -> None:
+        if key not in self._store_sum:
+            self.clear(key)
+        if isinstance(val, torch.Tensor):
+            val = val.item() # tensor -> num
+        self._store_sum[key] += val
+        self._store_total[key] += 1
+
+    def get(self, key: str) -> float:
+        total = max(self._store_total.get(key).item(), 1.0)
+        return (self._store_sum[key] / float(total)).item() or 0.0
+    
+    def clear(self, key: str) -> None:
+        self._store_sum[key] = torch.tensor(0.0, dtype=torch.float32)
+        self._store_total[key] = torch.tensor(0, dtype=torch.int32)
+    
+    def clear_all(self) -> None:
+        for key in self._store_sum:
+            self.clear(key)
+
+    def get_and_clear_all(self) -> Dict[str, float]:
+        metrics = {}
+        for key in self:
+            metrics[key] = self.get(key)
+            self.clear(key)
+        return metrics
+
+def load_embedder_and_tokenizer(name: str) -> Tuple[
+        transformers.PreTrainedModel, 
+        transformers.PreTrainedTokenizer
+]:
+    if name.startswith("nomic") or (name == "bert-base-uncased"):
+        from cde.lib.nomic_bert import NomicBertModel
+        if name.endswith("--from-scratch"):
+            name = name.replace("--from-scratch", "")
+            config = transformers.AutoConfig.from_pretrained(name, trust_remote_code=True)
+            model = NomicBertModel._from_config(config)
+        else:
+            model = NomicBertModel.from_pretrained(
+                name, add_pooling_layer=False
+            )
+        tokenizer = transformers.AutoTokenizer.from_pretrained(name)
+    elif name in ["gtr-base", "gtr_base"]:
+        model = transformers.AutoModel.from_pretrained(
+            "sentence-transformers/gtr-t5-base"
+        ).encoder
+        tokenizer = transformers.AutoTokenizer.from_pretrained(
+            "sentence-transformers/gtr-t5-base"
+        )
+    elif name == "pile-t5-base-encoder":
+        model = transformers.AutoModel.from_pretrained(
+            "EleutherAI/pile-t5-base"
+        ).encoder
+        tokenizer = transformers.AutoTokenizer.from_pretrained(
+            "EleutherAI/pile-t5-base"
+        )
+        tokenizer.pad_token = tokenizer.eos_token
+    elif name == "pile-t5-base-decoder":
+        model = transformers.AutoModel.from_pretrained(
+            "EleutherAI/pile-t5-base"
+        ).decoder
+        tokenizer = transformers.AutoTokenizer.from_pretrained(
+            "EleutherAI/pile-t5-base"
+        )
+        tokenizer.pad_token = tokenizer.eos_token
+    elif name.startswith("gpt2") or name.startswith("meta-llama") or ("Llama" in name):
+        model = transformers.AutoModelForCausalLM.from_pretrained(
+            name, 
+            # torch_dtype=torch.bfloat16,
+            attn_implementation="flash_attention_2" if torch.cuda.is_available() else "sdpa",
+            low_cpu_mem_usage=True,
+            # device_map="auto",
+        )
+        model.padding_side = "right"
+        tokenizer = transformers.AutoTokenizer.from_pretrained(name)
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.add_eos_token = True
+    elif "Modern" in name:
+        print("special loading for ModernBERT!")
+        # [1] needed for faster training
+        # model = transformers.AutoModel.from_pretrained(name, trust_remote_code=True, reference_compile=True)
+        # [2] needed for non-breaking inference
+        model = transformers.AutoModel.from_pretrained(name, trust_remote_code=True, reference_compile=False)
+        tokenizer = transformers.AutoTokenizer.from_pretrained(name)
+    else:
+        model = transformers.AutoModel.from_pretrained(name, trust_remote_code=True)
+        tokenizer = transformers.AutoTokenizer.from_pretrained(name)
+    return model, tokenizer
+
+
+def inputs_for_key(inputs: Dict[str, torch.Tensor], key: str):
+    key += "_"
+    return {k.replace(key, ""): v for k,v in inputs.items() if k.startswith(key)}
+
+
+def count_cpus() -> int:
+    try:
+        return len(os.sched_getaffinity(0)) 
+    except AttributeError:
+        return multiprocessing.cpu_count()
+
+
+def shuffle_batches(g: torch.Generator, list_of_tensors: List[torch.Tensor]) -> List[int]:
+    all_indices = []
+    for batch_tensor in tqdm_if_main_worker(list_of_tensors, colour="green", desc="Sampler shuffling per-batch"): 
+        rand_perm = torch.randperm(len(batch_tensor), generator=g)
+        batch_list = batch_tensor[rand_perm].tolist()
+        all_indices.extend(batch_list)
+    return all_indices
+
+
+# def shuffle_batches_multiproc(g: torch.Generator, list_of_tensors: List[torch.Tensor], num_processes: int = 8) -> List[int]:
+#     all_indices = []
+#     print(f"Shuffling {len(list_of_tensors)} tensors with {num_processes} workers.")
+#     pbar = tqdm_if_main_worker(list_of_tensors, colour="orange", desc=f"Sampler shuffling per-batch (nproc={num_processes})")
+#     pool = multiprocessing.Pool(processes=num_processes) 
+#     chunk_size = len(list_of_tensors) // num_processes
+#     chunks = [list_of_tensors[i:i + chunk_size] for i in range(0, len(list_of_tensors), chunk_size)]
+#     worker_func = functools.partial(shuffle_batches, g=g)
+#     results = pool.map(worker_func, chunks)
+#     all_indices = []
+#     for result in results:
+#         all_indices.extend(result)
+#         pbar.update()
+#     return all_indices
+
+
+def exit_if_running_or_finished_wandb(
+        project_name: str,
+        exp_group: str, exp_name: str
+    ) -> None:
+    print("Checking if experiment is already running...")
+    import wandb
+
+    api = wandb.Api()
+    running_runs = api.runs(
+        path="cde-0",
+        filters={
+            "display_name": exp_name,
+            "state": {"$regex": "Running|Finished"},
+            "config.exp_group": exp_group,
+        }  
+    )
+    print("Found", len(running_runs), f"runs with name {exp_name} and group {exp_group} in {project_name}.")
+
+    if len(running_runs) > 0:
+        print("Exiting because experiment is already running or completed.")
+        sys.exit(0)
+    
+
+HN_FILTER_TOKENIZER_MAP = {
+    "nomic": "nomic-ai/nomic-embed-text-v1",
+    "stella": "dunzhang/stella_en_400M_v5",
+    "sbert": "sentence-transformers/all-MiniLM-L6-v2",
+    "sentence_t5": "sentence-transformers/sentence-t5-base",
+    "gte": "Alibaba-NLP/gte-large-en-v1.5",
+}
+def load_hn_filter_tokenizer(tokenizer_name: str) -> Optional[transformers.PreTrainedTokenizer]:
+    if tokenizer_name in HN_FILTER_TOKENIZER_MAP:
+        return transformers.AutoTokenizer.from_pretrained(HN_FILTER_TOKENIZER_MAP[tokenizer_name])
+    else:
+        return None

model.py

@@ -0,0 +1,622 @@
+from typing import Optional
+
+import copy
+import torch
+import torch.nn as nn
+import transformers
+
+from cde.lib.dist import print0
+from cde.lib.tensor import mean_pool, mean_pool_3d, mean_pool_weighted, last_token_pool
+
+from cde.lib import load_embedder_and_tokenizer, ContextualModelConfig
+
+
+gpt_tokenizer = transformers.AutoTokenizer.from_pretrained("gpt2")
+
+def limit_layers(model: transformers.PreTrainedModel, n_layers: int) -> None:
+    if hasattr(model, 'transformer'):
+        if hasattr(model.transformer, 'h'):
+            # gpt2
+            model.transformer.h = model.transformer.h[:n_layers]
+        else:
+            model.transformer.layer = model.transformer.layer[:n_layers]
+    elif hasattr(model, 'encoder'):
+        if hasattr(model.encoder, 'layers'):
+            model.encoder.layers = model.encoder.layers[:n_layers]
+        else:
+            model.encoder.layer = model.encoder.layer[:n_layers]
+    else:
+        raise RuntimeError(f"unknown how to limit layers of model {type(model)}")
+
+
+def disable_dropout(model: torch.nn.Module):
+    dropout_modules = [m for m in model.modules() if isinstance(m, torch.nn.Dropout)]
+    for m in dropout_modules:
+        m.p = 0.0
+    print0(
+        f"Disabled {len(dropout_modules)} dropout modules from model type {type(model)}"
+    )
+
+
+def disable_causality(model: torch.nn.Module):
+    disabled_modules = 0
+    for m in model.modules():
+        if hasattr(m, "is_causal"):
+            m.is_causal = False
+            disabled_modules += 1
+    print0(
+        f"Set is_causal=False in {disabled_modules} modules from model type {type(model)}"
+    )
+
+
+class ContextualModelMixin(nn.Module):
+    @property
+    def num_corpus_tokens(self) -> int:
+        return self.transductive_corpus_size * self.transductive_tokens_per_document
+
+    def contextual_init(self):
+        self.n_soft_prompt = 8
+        self.prompt_projection = torch.nn.Sequential(
+            torch.nn.Linear(self.hidden_size, self.hidden_size),
+            torch.nn.ReLU(),
+            torch.nn.Linear(self.hidden_size, self.hidden_size * self.n_soft_prompt)
+        )
+        self.transductive_corpus_size = vars(self.config).get("transductive_corpus_size", 1)
+        self.transductive_tokens_per_document = vars(self.config).get("transductive_tokens_per_document", 1)
+        self.randomize_dataset_sequence_order = True
+        self.sequence_dropout_prob = vars(self.config).get("transductive_sequence_dropout_prob", 0.0)
+        if self.sequence_dropout_prob > 0.0:
+            self.sequence_dropout_null_embedding = torch.nn.Parameter(
+                torch.randn(self.hidden_size) * 0.01,
+                requires_grad = True
+            )       
+        self.output_projection = torch.nn.Sequential(
+            torch.nn.Linear(self.hidden_size, self.hidden_size),
+            torch.nn.ReLU(),
+            torch.nn.Linear(self.hidden_size, self.hidden_size)
+        )
+
+    def _prepare_dataset_embeddings(
+            self, 
+            input_ids: torch.Tensor, 
+            dataset_embeddings: torch.Tensor,
+            null_dataset_embedding: bool = False,
+        ) -> torch.Tensor:
+        if not isinstance(dataset_embeddings, torch.Tensor):
+            dataset_embeddings = torch.tensor(dataset_embeddings)
+
+        if len(dataset_embeddings.shape) == 2:
+            # Auto-expand for a batch.
+            dataset_embeddings = dataset_embeddings[None, :, :] # (b, d) -> (1, b, d)
+        dataset_embeddings = dataset_embeddings.to(input_ids.device)
+    
+        batch_size = input_ids.shape[0]
+        if (self.transductive_tokens_per_document > 1):
+            if self.training:
+                # Choose N random documents to fill our context window with.
+                # This logic is a little confusing but allows us to sample a
+                # different batch *per-document*
+                assert dataset_embeddings.shape[1] == self.transductive_tokens_per_document
+                R = torch.randint(
+                    low=0, 
+                    high=len(dataset_embeddings), 
+                    size=(batch_size, self.config.transductive_corpus_size), 
+                    device=dataset_embeddings.device
+                )
+                # TODO make this deterministic somehow for evaluation?
+                dataset_embeddings = dataset_embeddings[R].reshape((batch_size, self.num_corpus_tokens, self.hidden_size))
+            else:
+                dataset_embeddings = dataset_embeddings.reshape((1, self.num_corpus_tokens, self.hidden_size))
+                # print("reshaped to dataset_embeddings.shape =", dataset_embeddings.shape)
+
+        if dataset_embeddings.shape[1] > self.num_corpus_tokens:
+            # If too many dataset embeddings are passed in, just take the first N until
+            # we have the proper number.
+            dataset_embeddings = dataset_embeddings[:, :self.num_corpus_tokens, :]
+        
+        _, corpus_size, _hidden_size = dataset_embeddings.shape
+        if _ == 1:
+            # Auto-expand for a batch.
+            dataset_embeddings = dataset_embeddings.expand((batch_size, -1, -1))
+
+        if self.training and self.sequence_dropout_prob > 0.0:
+            sequence_dropout_mask = (
+                torch.rand((batch_size, corpus_size), device=dataset_embeddings.device) < self.sequence_dropout_prob
+            )
+            null_embeddings = self.sequence_dropout_null_embedding[None, None].expand(batch_size, corpus_size, -1)
+            dataset_embeddings = torch.where(
+                sequence_dropout_mask[..., None], null_embeddings, dataset_embeddings
+            )
+        elif null_dataset_embedding:
+            null_embeddings = self.sequence_dropout_null_embedding[None, None].expand(batch_size, corpus_size, -1)
+            dataset_embeddings = null_embeddings
+        
+        # backbone_max_seq_length = self.backbone.config.max_trained_positions
+        # assert batch_size + (2 * self.n_soft_prompt + corpus_size) <= backbone_max_seq_length, "too many hard negatives for backbone model"
+        soft_prompt = torch.ones((1, self.hidden_size), device=dataset_embeddings.device, dtype=dataset_embeddings.dtype)
+        soft_prompt = self.prompt_projection(soft_prompt).reshape((1, self.n_soft_prompt, self.hidden_size))
+        soft_prompt = soft_prompt.expand((len(dataset_embeddings), -1, -1)) # -> (b, 4+b, d) # soft_prompt.repeat((len(input_ids), 1, 1))  
+        soft_prompt = torch.cat((dataset_embeddings, soft_prompt), dim=1)
+        
+        return soft_prompt
+
+
+class BiEncoder(transformers.PreTrainedModel):
+    config_class = ContextualModelConfig
+    embedder: transformers.PreTrainedModel
+    def __init__(
+            self, 
+            config, #: transformers.PreTrainedConfig, 
+        ):
+        super().__init__(config=config)
+        embedder, _ = load_embedder_and_tokenizer(
+            config.embedder,
+        )
+
+        if config.limit_layers:
+            print0(f"Limiting layers to {config.limit_layers}")
+            limit_layers(embedder, config.limit_layers)
+    
+        self.embedder = embedder
+        # if ("t5" in embedder.config.model_type):
+        #     print0(f"using torch.compile() on embedder of type `{embedder.config.model_type}`")
+        #     self.embedder = torch.compile(self.embedder) 
+        self.hidden_size = self.embedder.config.hidden_size
+        # Allow pooling to multiple tokens per document
+        self.transductive_tokens_per_document = vars(self.config).get("transductive_tokens_per_document", 1)
+        self.mlp = torch.nn.Sequential(
+            torch.nn.Linear(self.hidden_size, self.hidden_size),
+            torch.nn.GELU(),
+            torch.nn.Linear(self.hidden_size, self.config.embedding_output_dim or self.hidden_size),
+        )
+        self.temp = config.logit_scale
+
+        if config.disable_dropout:
+            disable_dropout(self)
+        self.pooling_strategy = vars(config).get("pooling_strategy", "mean")
+
+    def forward(
+            self, 
+            input_ids: torch.Tensor,
+            attention_mask: torch.Tensor,
+            dataset_input_ids: Optional[torch.Tensor] = None,
+            dataset_attention_mask: Optional[torch.Tensor] = None,
+            token_type_ids = None,
+            output_hidden_states: bool = False,
+        ) -> torch.Tensor:
+        """
+        query_embedding (float torch.Tensor) - shape (batch_size, embedding_dim)
+        document_embeddings (float torch.Tensor) - shape (corpus_size, embedding_dim)
+            where the corpus_size >= batch_size and is structured like this:
+                [d1, d2, d3, hn1_1, hn1_2, hn2_1, hn2_2, hn3_1, hn3_2]
+                for a corpus with three documents and two hard negatives per document
+        """
+        del token_type_ids
+
+        outputs = (
+            self.embedder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+            ).last_hidden_state
+        )
+        if self.transductive_tokens_per_document > 1:
+            document_embeddings = None
+            batch_size, seq_length, output_dim = outputs.shape
+
+            if seq_length % self.transductive_tokens_per_document != 0:
+                # Pad to nearest multiple
+                n_extra_embeds = self.transductive_tokens_per_document - (seq_length % self.transductive_tokens_per_document)
+                outputs = torch.cat(
+                    (outputs, torch.zeros((batch_size, n_extra_embeds, output_dim), device=outputs.device)),
+                    dim=1
+                )
+                attention_mask = torch.cat(
+                    (attention_mask, torch.zeros((batch_size, n_extra_embeds), device=attention_mask.device)),
+                    dim=1
+                )
+                seq_length += n_extra_embeds
+                print(f"Added {n_extra_embeds} padding tokens to input_ids and attention_mask")
+            
+            # print("ftransductive_tokens_per_document {self.transductive_tokens_per_document} outputs.shape =", outputs.shape)
+
+            outputs = outputs.reshape(
+                (batch_size,  self.transductive_tokens_per_document, seq_length // self.transductive_tokens_per_document, output_dim)
+            )
+
+            attention_mask = attention_mask.reshape((batch_size, self.transductive_tokens_per_document, -1))
+            document_embeddings = mean_pool_3d(outputs, attention_mask)
+            
+            document_embeddings = document_embeddings.reshape((batch_size, self.transductive_tokens_per_document, output_dim))
+        else:
+            if self.pooling_strategy == "mean":
+                document_embeddings = mean_pool(outputs, attention_mask)
+            else:
+                document_embeddings = document_embeddings.max(dim=1)
+        output = self.mlp(document_embeddings)
+        # breakpoint()
+
+        if output_hidden_states:
+            return {
+                "hidden_states": outputs,
+                "pooled": output,
+            }
+        else:
+            return output
+
+
+class DatasetConditionedAutoregressive(transformers.PreTrainedModel, ContextualModelMixin):
+    def __init__(
+            self, 
+            config,
+            dataset_backbone: transformers.PreTrainedModel,
+            first_stage_hidden_size: int,
+        ):
+        super().__init__(config=config)
+        self.backbone = dataset_backbone
+        self.backbone_hidden_size = self.backbone.config.hidden_size
+        self.hidden_size = first_stage_hidden_size # Input token size
+        self.contextual_init()
+        disable_causality(self.backbone)
+        
+        self.pool_ignore_contextual_tokens = vars(self.config).get("pool_ignore_contextual_tokens", False)
+        self.pool_ignore_instruction_tokens = vars(self.config).get("pool_ignore_instruction_tokens", False)
+        self.pool_instruction_end_id = self.backbone.config.bos_token_id
+        
+        # Override contextual init
+        self.output_projection = torch.nn.Sequential(
+            torch.nn.Linear(self.backbone_hidden_size, self.backbone_hidden_size),
+            torch.nn.ReLU(),
+            torch.nn.Linear(self.backbone_hidden_size, self.backbone_hidden_size)
+        )
+        self._shift_rotary_embedding()
+                
+    @property
+    def num_corpus_tokens(self) -> int:
+        return self.config.transductive_corpus_size * self.transductive_tokens_per_document
+
+    @property
+    def corpus_token_ratio(self) -> float:
+        # How many tokens from the first stage make one token in the second
+        # stage?
+        return self.backbone_hidden_size / self.hidden_size
+    
+    def corpus_token_pad_size(self, n_tokens: int) -> int:
+        return self.hidden_size % self.backbone_hidden_size
+    
+    def _shift_rotary_embedding(self) -> None:
+        disable_transductive_rotary_embedding = vars(self.config).get("disable_transductive_rotary_embedding", True)
+        # TODO: Can we do this for LLAMA?
+        print0("Warning: Positional embedding disabling not implemented for LLAMA.")
+    
+    def forward(
+            self, 
+            input_ids: torch.Tensor,
+            attention_mask: torch.Tensor,
+            dataset_embeddings: torch.Tensor,
+            output_hidden_states: bool = False,
+            null_dataset_embedding: bool = False,
+        ) -> torch.Tensor:
+        soft_prompt = self._prepare_dataset_embeddings(
+            input_ids=input_ids,
+            dataset_embeddings=dataset_embeddings,
+            null_dataset_embedding=null_dataset_embedding,
+        )
+        
+        # Reshape for this model.
+        # print("[DatasetConditionedAutoregressive] 1 -> soft_prompt.shape =", soft_prompt.shape)
+        num_soft_elements = torch.prod(torch.tensor(soft_prompt.shape[1:])).item()
+        soft_prompt = soft_prompt.reshape((soft_prompt.shape[0], num_soft_elements))
+        num_padding_elements = self.backbone_hidden_size - (num_soft_elements % self.backbone_hidden_size)
+        padding = torch.ones((soft_prompt.shape[0], num_padding_elements), device=soft_prompt.device)
+        soft_prompt = torch.cat((soft_prompt, padding), dim=1)
+        soft_prompt = soft_prompt.reshape(
+            (soft_prompt.shape[0], -1, self.backbone_hidden_size)
+        )
+        # print("[DatasetConditionedAutoregressive] 2 -> soft_prompt.shape =", soft_prompt.shape)
+
+        backbone_attention_mask = torch.ones(
+            soft_prompt.shape[0:2],
+            dtype=torch.long,
+            device=soft_prompt.device,
+        )
+        token_embeddings = self.backbone.get_input_embeddings()
+        inputs_embeds = token_embeddings(input_ids) # (b, s) -> (b, s, d)
+        # print("[2] inputs_embeds.shape =", inputs_embeds.shape)
+        inputs_embeds = torch.cat((soft_prompt, inputs_embeds), dim=1) # (v, 4+b+s, d)
+        # print("[3.a] inputs_embeds.shape =", inputs_embeds.shape)
+        input_attention_mask = torch.cat((backbone_attention_mask, attention_mask), dim=1)
+        # print("[3.b] attention_mask.shape =", attention_mask.shape)
+
+        output = self.backbone(
+            inputs_embeds=inputs_embeds,
+            attention_mask=input_attention_mask,
+            output_hidden_states=True,
+        ) # (1, 4 + b + s, d)
+        # trim soft prompt
+        output_vectors = output.hidden_states[-1]
+        n_soft_prompt_tokens = soft_prompt.shape[1]
+
+        if self.pool_ignore_instruction_tokens:
+            # Denote the end of an instruction with an extra BOS token.
+            # This is a bit arcane but relies on the fact that there will be a BOS token after the
+            # instruction, but also there may or may not be a BOS token at the beginning.
+            instruction_end_idx = (
+                (input_ids == self.pool_instruction_end_id) & 
+                attention_mask &
+                (torch.arange(input_ids.shape[1], device=input_ids.device)[None, :] > 0)
+            ).int().argmax(1)
+            is_instruction_token_mask = (
+                torch.arange(input_ids.shape[1], device=input_ids.device)[None, :] <= instruction_end_idx[:, None]
+            )
+            # catch edge case where there is no instruction
+            is_instruction_token_mask = is_instruction_token_mask.where(
+                (instruction_end_idx > 0)[:, None], torch.zeros_like(is_instruction_token_mask)
+            )
+            input_attention_mask = torch.cat((
+                backbone_attention_mask, 
+                attention_mask & ~is_instruction_token_mask), dim=1
+            )
+
+        output_attention_mask = input_attention_mask
+        if self.pool_ignore_contextual_tokens:
+            output_vectors = output_vectors[:, n_soft_prompt_tokens:, :]
+            output_attention_mask = output_attention_mask[:, n_soft_prompt_tokens:]
+
+        # Take last token position
+        if vars(self.config).get("pooling_strategy") == "last_token":
+            output_pooled = last_token_pool(output_vectors, output_attention_mask)
+        elif vars(self.config).get("pooling_strategy") == "mean":
+            output_pooled = mean_pool(output_vectors, output_attention_mask)
+        else:
+            output_pooled = mean_pool_weighted(output_vectors, output_attention_mask)
+
+        # average with original vectors
+        output = self.output_projection(output_pooled) # (b, 2d) -> (b, d)
+
+        if output_hidden_states:
+            return {
+                "hidden_states": output_vectors,
+                "pooled": output,
+            }
+        else:
+            return output
+
+
+class DatasetConditionedBiencoder(transformers.PreTrainedModel, ContextualModelMixin):
+    def __init__(
+            self, 
+            config,
+            dataset_backbone: transformers.PreTrainedModel,
+        ):
+        super().__init__(config=config)
+        self.backbone = dataset_backbone
+        self.hidden_size = self.backbone.config.hidden_size
+        self.hidden_size = dataset_backbone.config.hidden_size
+        self.contextual_init()
+        self._shift_rotary_embedding()
+
+        self.pool_ignore_contextual_tokens = vars(self.config).get("pool_ignore_contextual_tokens", True)
+        self.pool_ignore_instruction_tokens = vars(self.config).get("pool_ignore_instruction_tokens", False)
+
+        tokenizer = transformers.AutoTokenizer.from_pretrained(self.config.embedder)
+        self.pool_instruction_end_id = tokenizer.encode(": ", add_special_tokens=False)[0] # Hardcoded for colon-ending prefixes.
+                
+    @property
+    def num_corpus_tokens(self) -> int:
+        return self.config.transductive_corpus_size * self.transductive_tokens_per_document
+    
+    def _shift_rotary_embedding(self) -> None:
+        disable_transductive_rotary_embedding = vars(self.config).get("disable_transductive_rotary_embedding", True)
+        if self.backbone.config.model_type.startswith("nomic") and disable_transductive_rotary_embedding:
+            # We only want to apply positional embeddings to the
+            # *text* portion of the backbone network.
+            self.backbone.config.rotary_start_pos = 0.0
+            rotary_disabled = 0
+
+            rotary_start_pos = self.num_corpus_tokens
+            for module in self.backbone.modules():
+                if hasattr(module, "rotary_emb_dim"):
+                    module.rotary_start_pos = rotary_start_pos
+                    rotary_disabled += 1
+            print0(f"modified {rotary_disabled} rotary modules – set rotary_start_pos to {rotary_start_pos}")
+    
+    def forward(
+            self, 
+            input_ids: torch.Tensor,
+            attention_mask: torch.Tensor,
+            dataset_embeddings: torch.Tensor,
+            output_hidden_states: bool = False,
+            null_dataset_embedding: bool = False,
+        ) -> torch.Tensor:
+        soft_prompt = self._prepare_dataset_embeddings(
+            input_ids=input_ids,
+            dataset_embeddings=dataset_embeddings,
+            null_dataset_embedding=null_dataset_embedding,
+        )
+        backbone_attention_mask = torch.ones(
+            soft_prompt.shape[0:2],
+            dtype=torch.long,
+            device=soft_prompt.device,
+        )
+        inputs_embeds = self.backbone.embeddings(input_ids) # (b, s) -> (b, s, d)
+        inputs_embeds = torch.cat((soft_prompt, inputs_embeds), dim=1) # (v, 4+b+s, d)
+        input_attention_mask = torch.cat((backbone_attention_mask, attention_mask), dim=1)
+        output = self.backbone(
+            inputs_embeds=inputs_embeds,
+            attention_mask=input_attention_mask,
+        ) # (1, 4 + b + s, d)
+        # trim soft prompt
+        output_vectors = output.last_hidden_state
+
+        # use only these tokens
+        n_soft_prompt_tokens = soft_prompt.shape[1]
+
+        if self.pool_ignore_instruction_tokens:
+            # Denote the end of an instruction with an extra BOS token.
+            # This is a bit arcane but relies on the fact that there will be a BOS token after the
+            # instruction, but also there may or may not be a BOS token at the beginning.
+            instruction_end_idx = (
+                (input_ids == self.pool_instruction_end_id) & 
+                attention_mask &
+                (torch.arange(input_ids.shape[1], device=input_ids.device)[None, :] > 0)
+            ).int().argmax(1)
+            is_instruction_token_mask = (
+                torch.arange(input_ids.shape[1], device=input_ids.device)[None, :] <= instruction_end_idx[:, None]
+            )
+            # catch edge case where there is no instruction
+            is_instruction_token_mask = is_instruction_token_mask.where(
+                (instruction_end_idx > 0)[:, None], torch.zeros_like(is_instruction_token_mask)
+            )
+            output_attention_mask = torch.cat((backbone_attention_mask, attention_mask & ~is_instruction_token_mask), dim=1)
+        else:
+            output_attention_mask = input_attention_mask
+
+        if self.pool_ignore_contextual_tokens:
+            output_vectors = output_vectors[:, n_soft_prompt_tokens:, :]
+            output_attention_mask = output_attention_mask[:, n_soft_prompt_tokens:]
+        output_pooled = mean_pool(output_vectors, output_attention_mask)
+        # average with original vectors
+        output = self.output_projection(output_pooled) + output_pooled # (b, d) -> (b, d) / with residual connection
+
+        if output_hidden_states:
+            return {
+                "hidden_states": output_vectors,
+                "pooled": output,
+            }
+        else:
+            return output
+
+
+class DatasetPrefixBiencoder(transformers.PreTrainedModel, ContextualModelMixin):
+    def __init__(
+            self, 
+            config, #: transformers.PreTrainedConfig, 
+            embedder: transformers.PreTrainedModel, 
+        ):
+        super().__init__(config=config)
+        self.embedder = embedder
+        self.hidden_size = self.embedder.config.hidden_size
+        self.contextual_init()
+    
+    def forward(
+            self, 
+            input_ids: torch.Tensor,
+            attention_mask: torch.Tensor,
+            dataset_input_ids: torch.Tensor,
+            dataset_attention_mask: torch.Tensor,
+            output_hidden_states: bool = False,
+        ) -> torch.Tensor:
+        R = torch.randint(low=0, high=len(dataset_input_ids), size=(len(input_ids),), device=dataset_input_ids.device)
+        
+        dataset_input_ids = dataset_input_ids[R]
+        input_ids = torch.cat((dataset_input_ids, input_ids), dim=1)
+
+        dataset_attention_mask = torch.ones_like(dataset_attention_mask, device=dataset_attention_mask.device)
+        input_attention_mask = torch.cat((dataset_attention_mask, attention_mask), dim=1)
+        output_attention_mask = torch.cat(
+            (torch.zeros_like(dataset_input_ids), attention_mask), dim=1
+        )
+
+        output = self.embedder(
+            input_ids=input_ids,
+            attention_mask=input_attention_mask,
+        ) 
+        
+        output_vectors = output.last_hidden_state
+        output_pooled = mean_pool(output_vectors, output_attention_mask)
+        output = self.output_projection(output_pooled) # (b, 2d) -> (b, d)
+
+        if output_hidden_states:
+            S_d = dataset_attention_mask.shape[1]
+            output_vectors = output_vectors[:, S_d:, :]
+            return {
+                "hidden_states": output_vectors,
+                "pooled": output,
+            }
+        else:
+            return output
+
+
+class ContextualDocumentEmbeddingTransformer(transformers.PreTrainedModel):
+    config_class = ContextualModelConfig
+    embedder: transformers.PreTrainedModel
+    dataset_backbone: transformers.PreTrainedModel
+    def __init__(
+            self, 
+            config,
+        ):
+        super().__init__(config=config)
+        dataset_backbone, _ = load_embedder_and_tokenizer(
+            vars(config).get("dataset_backbone") or config.embedder
+        )
+
+        if config.limit_layers:
+            print0(f"Limiting layers to {config.limit_layers}")
+            limit_layers(dataset_backbone, config.limit_layers)
+        
+        biencoder_config = copy.deepcopy(config)
+        biencoder_config.embedding_output_dim = None
+        biencoder_config.limit_layers = vars(self.config).get("limit_layers_first_stage", None)
+        self.first_stage_model = BiEncoder(
+            config=biencoder_config,
+        )
+
+        if vars(config).get("autoregressive_backbone", False):
+            self.second_stage_model = DatasetConditionedAutoregressive(
+                config=config,
+                dataset_backbone=dataset_backbone,
+                first_stage_hidden_size=self.first_stage_model.hidden_size,
+            )
+        else:
+            self.second_stage_model = DatasetConditionedBiencoder(
+                config=config,
+                dataset_backbone=dataset_backbone
+            )
+        
+        self.temp = config.logit_scale
+        if config.disable_dropout:
+            disable_dropout(self)
+        
+        transductive_tie_token_embeddings = vars(self.config).get("transductive_tie_token_embeddings", False)
+        if transductive_tie_token_embeddings:
+            self.second_stage_model.backbone.embeddings.word_embeddings.weight = (
+                self.first_stage_model.embedder.embeddings.word_embeddings.weight
+            )
+
+    def forward(
+            self, 
+            input_ids: torch.Tensor,
+            attention_mask: torch.Tensor,
+            dataset_input_ids: Optional[torch.Tensor],
+            dataset_attention_mask: Optional[torch.Tensor],
+            output_hidden_states: bool = False,
+        ) -> torch.Tensor:
+        """
+        input_ids (long torch.Tensor) – ids of input tokens
+        attention_mask (bool torch.Tensor)
+        """
+        dataset_embeddings = self.first_stage_model(
+            input_ids=dataset_input_ids, 
+            attention_mask=dataset_attention_mask
+        )
+        return self.second_stage_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            dataset_embeddings=dataset_embeddings,
+            output_hidden_states=output_hidden_states,
+        )
+
+
+
+def get_model_class(name: str):
+    if name in 'transductive': 
+        return ContextualDocumentEmbeddingTransformer
+    elif name == 'biencoder':
+        return BiEncoder
+    elif name == "biencoder_plus_plus":
+        from cde.model_extra import BiEncoderPlusPlus
+        return BiEncoderPlusPlus
+    elif name == "dataset_prefix_biencoder":
+        return DatasetPrefixBiencoder
+    else:
+        raise ValueError(f'unknown model cls {name}')

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e7cca261c510de07c012f3019366f1b6c5720761b6966b0388faea6e70398983
+size 1124594680