Upload ContextualDocumentEmbeddingTransformer

config.json

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

model.py

@@ -1,439 +1,17 @@
-###################################################################################################
-###################################################################################################
-###################################################################################################
+from typing import Optional
 
-import collections
-import logging
-
-import json
-import math
-import os
-import re
-from collections import OrderedDict
-from functools import partial
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-
-
-
-########################################################
-########################################################
-########################################################
-########################################################
-
-
-from typing import Callable, Optional, Tuple
 import copy
-import math
-import multiprocessing
-import os
-
 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
 
-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 load_embedder_and_tokenizer(name: str) -> Tuple[
-        transformers.PreTrainedModel, 
-        transformers.PreTrainedTokenizer
-]:
-    print("Loading model:", name)
-    if name.startswith("nomic") or (name == "bert-base-uncased"):
-        model = ContextualNomicBertForPreTraining.from_pretrained(name, trust_remote_code=True).bert
-        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",
-            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
-    else:
-        model = transformers.AutoModel.from_pretrained(name, trust_remote_code=True)
-        tokenizer = transformers.AutoTokenizer.from_pretrained(name)
-
-        # if use_bettertransformer:
-        #     from optimum.bettertransformer import BetterTransformer
-        #     model = BetterTransformer.transform(model)
-    return model, tokenizer
-
-
-def get_world_size() -> int:
-    try:
-        return torch.distributed.get_world_size()
-    except (RuntimeError, ValueError):
-        return 1
-
-
-def get_rank() -> int:
-    try:
-        return torch.distributed.get_rank()
-    except (RuntimeError, ValueError):
-        return 0
-    
-def gather(t: torch.Tensor) -> torch.Tensor:
-    # torch.distributed.nn.all_gather scales by world size since the reduce op is SUM
-    # https://github.com/pytorch/pytorch/issues/58005
-    # only should use torch.distributed.nn.all_gather if we implement a `local_loss`
-    # like: https://github.com/mlfoundations/open_clip/issues/616
-    world_size = get_world_size()
-    if world_size == 1:
-        return t
-
-    if t.ndim == 0:
-        t = t.unsqueeze(0)
-
-    gathered = [torch.empty_like(t) for _ in range(world_size)]
-    torch.distributed.all_gather(gathered, t)
-    gathered[get_rank()] = t
-    return torch.cat(gathered, dim=0)
-
-
-def gather_sum(t: torch.Tensor) -> torch.Tensor:
-    # torch.distributed.nn.all_gather scales by world size since the reduce op is SUM
-    # https://github.com/pytorch/pytorch/issues/58005
-    # only should use torch.distributed.nn.all_gather if we implement a `local_loss`
-    # like: https://github.com/mlfoundations/open_clip/issues/616
-    world_size = get_world_size()
-    if world_size == 1:
-        return t
-
-    if t.ndim == 0:
-        t = t.unsqueeze(0)
-
-    gathered = [torch.empty_like(t) for _ in range(world_size)]
-    torch.distributed.all_gather(gathered, t)
-    gathered = torch.stack(gathered, dim=0)
-    return gathered.sum(dim=0) # Sum across workers
+from cde.lib import load_embedder_and_tokenizer, ContextualModelConfig
 
 
-def get_num_proc() -> int:
-    world_size: int = get_world_size()
-    try:
-        # os.sched_getaffinity respects schedulers, unlike cpu_count(), but it's only available
-        # on some Unix platforms, so we support both!
-        return len(os.sched_getaffinity(0)) // world_size  # type: ignore[attr-defined]
-    except AttributeError:
-        return multiprocessing.cpu_count() // world_size
-
-
-def torch_main_worker_finish_first(func: Callable):
-    def wrapper(*args, **kwargs):
-        # Get local rank (need to support non-DDP).
-        try:
-            local_rank = torch.distributed.get_rank()
-            ddp_enabled = True
-        except (RuntimeError, ValueError):
-            local_rank = -1
-            ddp_enabled = False
-        is_main_worker = local_rank <= 0
-        # Run on main worker first.
-        if is_main_worker:
-            result = func(*args, **kwargs)
-        # Then everyone waits.
-        if ddp_enabled:
-            torch.distributed.barrier()
-        # Run on other workers now.
-        if not is_main_worker:
-            result = func(*args, **kwargs)
-        # Now everyone waits again.
-        if ddp_enabled:
-            torch.distributed.barrier()
-        return result
-
-    return wrapper
-
-
-def print0(*args, **kwargs) -> None:
-    if get_rank() == 0:
-        print(*args, **kwargs)
-
-
-def verify_ddp_weights_equal(model: torch.nn.Module, atol: float = 1e-5) -> None:
-    if hasattr(model, "module"):
-        model = model.module
-    
-    world_size = get_world_size()
-
-    if world_size > 8:
-        print0(f"[verify_ddp_weights_equal] Skipping with world_size={world_size} ⚠️")
-        return
-
-    for name, param in model.named_parameters():
-        if param is None: continue
-        if param.grad is None: 
-            print0(f"[verify_ddp_weights_equal] Skipping param [{name}] with no grad")
-            continue
-        gathered_param = gather(param).reshape((world_size, -1))
-        absolute_diffs = (gathered_param[None, 0, :] - gathered_param).abs()
-        rank_params_eq = (absolute_diffs < atol).all()
-        assert rank_params_eq, f"❌ param [{name}] not equal - got max_absolute_diff={absolute_diffs.max()}"
-        ###################################################################################################################
-        gathered_param_grad = gather(param.grad).reshape((world_size, -1))
-        absolute_grad_diffs = (gathered_param_grad[None, 0, :] - gathered_param_grad).abs()
-        rank_grad_params_eq = (absolute_grad_diffs < atol).all()
-        assert rank_grad_params_eq, f"❌ param [{name}] grad not equal - got max_absolute_diff={absolute_grad_diffs.max()}"
-        ###################################################################################################################
-        
-    
-    print0("[verify_ddp_weights_equal] Verified DDP parameter correctness ✅")
-    
-
-
-def mean_pool_3d(
-    hidden_states: torch.Tensor, attention_mask: torch.Tensor
-) -> torch.Tensor:
-    B, T, S, D = hidden_states.shape
-    unmasked_outputs = hidden_states * attention_mask[..., None]
-    pooled_outputs = unmasked_outputs.sum(dim=2) / (attention_mask.sum(dim=2)[..., None] + 1e-9)
-
-    # fix for gradient flow: fill empty rows with the mean of the rest of the sequence
-    sequence_means = (
-        hidden_states.reshape((B, S * T, D))
-            .mean(dim=1, keepdim=True)
-            .expand(-1, T, -1)
-    )
-    pooled_outputs = pooled_outputs.where(
-        (attention_mask.sum(dim=2)[..., None] > 0), 
-        sequence_means
-    )
-    assert pooled_outputs.shape == (B, T, D)
-
-    return pooled_outputs
-
-def mean_pool(
-    hidden_states: torch.Tensor, attention_mask: torch.Tensor
-) -> torch.Tensor:
-    B, _S, D = hidden_states.shape
-    unmasked_outputs = hidden_states * attention_mask[..., None]
-    pooled_outputs = unmasked_outputs.sum(dim=1) / (attention_mask.sum(dim=1)[:, None] + 1e-20)
-    
-    assert pooled_outputs.shape == (B, D)
-    return pooled_outputs
-
-
-def mean_pool_weighted(
-    hidden_states: torch.Tensor, attention_mask: torch.Tensor
-) -> torch.Tensor:
-    B, _S, D = hidden_states.shape
-    attention_mask *= attention_mask.cumsum(dim=1) # [0,1,1,1,0,0] -> [0,1,2,3,0,0]
-    s = torch.sum(hidden_states * attention_mask.unsqueeze(-1).float(), dim=1)
-    d = attention_mask.sum(dim=1, keepdim=True).float()
-    return s / d
-
-
-def slice_sparse_tensor_rows(t: torch.sparse.Tensor, min_row: int, max_row: int) -> torch.sparse.Tensor:
-    assert min_row < max_row, f"can't slice from row {min_row} to {max_row}"
-    t = t.coalesce()
-    row_idxs = t.indices()[0]
-    index_mask = (min_row <= row_idxs) & (row_idxs < max_row)
-
-    num_rows = (max_row - min_row)
-    num_cols = t.shape[1]
-
-    idxs = t.indices()[:, index_mask]
-    vals = t.values()[index_mask]
-    return torch.sparse_coo_tensor(idxs, vals, size=(num_rows, num_cols)).coalesce()
-
-
-def slice_tensor_rows(t: torch.Tensor, min_row: int, max_row: int) -> torch.Tensor:
-    if t.is_sparse:
-        return slice_sparse_tensor_rows(t=t, min_row=min_row, max_row=max_row)
-    else:
-        return t[min_row:max_row]
-
-
-@torch.no_grad
-def maxsim(
-    X: torch.Tensor, y: torch.Tensor, 
-    maximize: bool, chunk_size: int = 8_000,
-    debug_mem_usage: bool = False) -> torch.Tensor:
-    device = X.device
-    n_samples = X.shape[0]
-
-    max_sim_v = torch.zeros(n_samples, device=device, dtype=X.dtype)
-    max_sim_i = torch.zeros(n_samples, device=device, dtype=torch.int64)
-
-    # TODO: Implement faster max (without going to dense tensors).
-    # TODO: Use multiple GPUs.
-    rank = get_rank()
-    world_size = get_world_size()
-
-    worker_worklist_size = int(math.ceil(n_samples / world_size))
-    splits_start_idx = worker_worklist_size * rank
-    splits_end_idx = worker_worklist_size * (rank + 1)
-
-    for i in range(splits_start_idx, splits_end_idx, chunk_size):
-        start, end = i, min(i + chunk_size, n_samples)
-        sub_x = slice_tensor_rows(X, start, end)
-        if debug_mem_usage: print(f"[maxsim] step {i} cuda mem free/total = {torch.cuda.mem_get_info()}")
-        if debug_mem_usage: print("[maxsim] sub_x.shape:", sub_x.shape, "//", "y.shape:", y.shape)
-        sub_sim = sub_x @ y # TODO – Implement sparse max here to save mem!
-        sub_sim = sub_sim
-        if maximize:
-            sub_max_sim_v, sub_max_sim_i = sub_sim.to_dense().max(dim=-1)
-        else:
-            sub_max_sim_v, sub_max_sim_i = sub_sim.to_dense().min(dim=-1)
-        del sub_sim
-        del sub_x
-        torch.cuda.empty_cache() # needs to happen after maxsim for some reason.
-        max_sim_v[start: end] = sub_max_sim_v
-        max_sim_i[start: end] = sub_max_sim_i
-    
-    # gather
-    max_sim_v = gather_sum(max_sim_v)
-    max_sim_i = gather_sum(max_sim_i)
-    k = y.shape[1]
-
-    assert max_sim_v.shape == (n_samples,)
-    assert max_sim_i.shape == (n_samples,)
-    assert max_sim_i.min() >= 0
-    assert max_sim_i.max() <= k
-
-    return max_sim_v, max_sim_i
-
-
-def forward_batched(
-        model: torch.nn.Module,
-        input_ids: torch.Tensor,
-        attention_mask: torch.Tensor,
-        batch_size: int,
-        dataset_input_ids: Optional[torch.Tensor] = None,
-        dataset_attention_mask: Optional[torch.Tensor] = None,
-        **second_stage_model_kwargs,
-) -> torch.Tensor:
-    if hasattr(model, "module"):
-        model = model.module
-    
-    if hasattr(model, "first_stage_model"):
-        # Support pooling over 3D dataset_input_ids inputs.
-        if len(dataset_input_ids.shape) == 2:
-            dataset_input_ids = dataset_input_ids[None]
-            dataset_attention_mask = dataset_attention_mask[None]
-
-        dataset_embeddings = []
-        for j in range(len(dataset_input_ids)):
-            i = 0
-            dataset_embeddings_batch = []
-            while i < dataset_input_ids.shape[1]:
-                dataset_embeddings_batch.append(
-                    model.first_stage_model(
-                        input_ids=dataset_input_ids[j][i:i+batch_size],
-                        attention_mask=dataset_attention_mask[j][i:i+batch_size],
-                    )
-                )
-                i += batch_size
-            dataset_embeddings.append(
-                torch.cat(dataset_embeddings_batch, dim=0)
-            )
-       
-        # Automatically pool over 3D dataset_input_ids.
-        dataset_embeddings = torch.stack(dataset_embeddings, dim=0).mean(dim=0)
-
-        j = 0
-        outputs = []
-        while j < len(input_ids):
-            outputs.append(
-                model.second_stage_model(
-                    input_ids=input_ids[j:j+batch_size],
-                    attention_mask=attention_mask[j:j+batch_size],
-                    dataset_embeddings=dataset_embeddings,
-                    **second_stage_model_kwargs,
-                )
-            )
-            j += batch_size
-        return torch.cat(outputs, dim=0)
-
-    else:
-        i = 0
-        outputs = []
-        while i < len(input_ids):
-            outputs.append(
-                model(
-                    input_ids=input_ids[i:i+batch_size],
-                    attention_mask=attention_mask[i:i+batch_size],
-                    **second_stage_model_kwargs,
-                )
-            )
-            i += batch_size
-        return torch.cat(outputs, dim=0)
-
-
-def last_token_pool(hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
-    # https://github.com/ContextualAI/gritlm/blob/main/gritlm/gritlm.py#L190
-    b, n, d = hidden_state.size()
-    # Get the last `1` in the attention mask of each item
-    # Often it is just `gather_indices = torch.argmin(attention_mask, 1, keepdim=False) - 1`
-    # except when 1) There's all 1's 2) There's 0's before the 1's
-    reversed_mask = torch.flip(attention_mask, dims=(1,))
-    argmax_reverse = torch.argmax(reversed_mask, dim=1, keepdim=False)
-    gather_indices = attention_mask.size(1) - argmax_reverse - 1
-    # If there are empty sequences, where the index would become -1 it will crash so set them to 0
-    gather_indices = torch.clamp(gather_indices, min=0)
-    # Turn indices from shape [b] -> [b, 1, d]
-    gather_indices = gather_indices.unsqueeze(-1).repeat(1, d)
-    gather_indices = gather_indices.unsqueeze(1)
-    assert gather_indices.shape == (b, 1, d)
-    # Gather along the seq len: [b, n, d] -> [b, d]
-    # Actually no need for the attention mask as we gather the last token where attn_mask=1 but
-    # as some indices (which shouldn't be attended to) may be 0 due to clamp, use mask to ignore them again
-    input_mask_expanded = attention_mask.unsqueeze(-1).expand((b, n, d)).float()
-    return torch.gather(hidden_state * input_mask_expanded, 1, gather_indices).squeeze(dim=1)
-
-def print0(*args, **kwargs) -> None:
-    if get_rank() == 0:
-        print(*args, **kwargs)
-
+gpt_tokenizer = transformers.AutoTokenizer.from_pretrained("gpt2")
 
 def limit_layers(model: transformers.PreTrainedModel, n_layers: int) -> None:
     if hasattr(model, 'transformer'):
@@ -449,7 +27,6 @@ def limit_layers(model: transformers.PreTrainedModel, n_layers: int) -> None:
             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):
@@ -501,7 +78,8 @@ class ContextualModelMixin(nn.Module):
 
     def _prepare_dataset_embeddings(
             self, 
-            input_ids: torch.Tensor, dataset_embeddings: torch.Tensor,
+            input_ids: torch.Tensor, 
+            dataset_embeddings: torch.Tensor,
             null_dataset_embedding: bool = False,
         ) -> torch.Tensor:
         if not isinstance(dataset_embeddings, torch.Tensor):
@@ -511,9 +89,6 @@ class ContextualModelMixin(nn.Module):
             # Auto-expand for a batch.
             dataset_embeddings = dataset_embeddings[None, :, :] # (b, d) -> (1, b, d)
         dataset_embeddings = dataset_embeddings.to(input_ids.device)
-
-        if len(dataset_embeddings.shape) < 3:
-            raise ValueError(f"dataset_embeddings must have at least 3 dimensions, got {dataset_embeddings.shape}")
     
         batch_size = input_ids.shape[0]
         if (self.transductive_tokens_per_document > 1):
@@ -532,11 +107,9 @@ class ContextualModelMixin(nn.Module):
                 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:
-            raise ValueError(f"dataset_embeddings must have at least {self.num_corpus_tokens} tokens, got {dataset_embeddings.shape[1]}")
-        elif dataset_embeddings.shape[1] > self.num_corpus_tokens:
+        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, :]
@@ -564,24 +137,12 @@ class ContextualModelMixin(nn.Module):
         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)
-
-        # print(f"[ContextualModelMixin] soft_prompt.shape = {soft_prompt.shape}")
-
-        if self.training and self.randomize_dataset_sequence_order:
-            randomized_order = torch.stack(
-                [
-                    torch.cat(
-                        (
-                            torch.randperm(corpus_size, device=soft_prompt.device), 
-                            torch.arange(self.n_soft_prompt, device=soft_prompt.device) + corpus_size
-                        ), dim=0) 
-                        for _ in range(batch_size)])
-            randomized_order = randomized_order.to(soft_prompt.device)
-            soft_prompt = soft_prompt.gather(1, randomized_order[..., None].expand_as(soft_prompt))
         
         return soft_prompt
 
+
 class BiEncoder(transformers.PreTrainedModel):
+    config_class = ContextualModelConfig
     embedder: transformers.PreTrainedModel
     def __init__(
             self, 
@@ -638,7 +199,6 @@ class BiEncoder(transformers.PreTrainedModel):
                 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
@@ -673,6 +233,7 @@ class BiEncoder(transformers.PreTrainedModel):
             else:
                 document_embeddings = document_embeddings.max(dim=1)
         output = self.mlp(document_embeddings)
+        # breakpoint()
 
         if output_hidden_states:
             return {
@@ -697,10 +258,9 @@ class DatasetConditionedAutoregressive(transformers.PreTrainedModel, ContextualM
         self.contextual_init()
         disable_causality(self.backbone)
         
-        self.input_ln = torch.nn.LayerNorm(
-            self.backbone_hidden_size, 
-            eps=1e-5
-        )
+        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(
@@ -726,7 +286,7 @@ class DatasetConditionedAutoregressive(transformers.PreTrainedModel, ContextualM
     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?
-        print("Warning: Positional embedding disabling not implemented for LLAMA.")
+        print0("Warning: Positional embedding disabling not implemented for LLAMA.")
     
     def forward(
             self, 
@@ -752,7 +312,6 @@ class DatasetConditionedAutoregressive(transformers.PreTrainedModel, ContextualM
         soft_prompt = soft_prompt.reshape(
             (soft_prompt.shape[0], -1, self.backbone_hidden_size)
         )
-        soft_prompt = self.input_ln(soft_prompt)
         # print("[DatasetConditionedAutoregressive] 2 -> soft_prompt.shape =", soft_prompt.shape)
 
         backbone_attention_mask = torch.ones(
@@ -774,11 +333,34 @@ class DatasetConditionedAutoregressive(transformers.PreTrainedModel, ContextualM
             output_hidden_states=True,
         ) # (1, 4 + b + s, d)
         # trim soft prompt
-        last_hidden_state = output.hidden_states[-1]
+        output_vectors = output.hidden_states[-1]
         n_soft_prompt_tokens = soft_prompt.shape[1]
 
-        output_vectors = last_hidden_state[:, n_soft_prompt_tokens:, :]
-        output_attention_mask = input_attention_mask[:, n_soft_prompt_tokens:]
+        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":
@@ -789,7 +371,6 @@ class DatasetConditionedAutoregressive(transformers.PreTrainedModel, ContextualM
             output_pooled = mean_pool_weighted(output_vectors, output_attention_mask)
 
         # average with original vectors
-        # TODO: Argparse for pooling strategy.
         output = self.output_projection(output_pooled) # (b, 2d) -> (b, d)
 
         if output_hidden_states:
@@ -801,7 +382,6 @@ class DatasetConditionedAutoregressive(transformers.PreTrainedModel, ContextualM
             return output
 
 
-
 class DatasetConditionedBiencoder(transformers.PreTrainedModel, ContextualModelMixin):
     def __init__(
             self, 
@@ -967,7 +547,7 @@ class ContextualDocumentEmbeddingTransformer(transformers.PreTrainedModel):
         ):
         super().__init__(config=config)
         dataset_backbone, _ = load_embedder_and_tokenizer(
-            vars(config).get("dataset_backbone", config.embedder)
+            vars(config).get("dataset_backbone") or config.embedder
         )
 
         if config.limit_layers:
@@ -1012,7 +592,7 @@ class ContextualDocumentEmbeddingTransformer(transformers.PreTrainedModel):
             output_hidden_states: bool = False,
         ) -> torch.Tensor:
         """
-        input_ids (long torch.Tensor) – ids of input tokens
+        input_ids (long torch.Tensor) – ids of input tokens
         attention_mask (bool torch.Tensor)
         """
         dataset_embeddings = self.first_stage_model(
@@ -1026,11 +606,16 @@ class ContextualDocumentEmbeddingTransformer(transformers.PreTrainedModel):
             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:

model.safetensors

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