re-add source w residual

model.py

@@ -1,17 +1,439 @@
-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
 
-from cde.lib import load_embedder_and_tokenizer, ContextualModelConfig
+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
 
 
-gpt_tokenizer = transformers.AutoTokenizer.from_pretrained("gpt2")
+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)
+
 
 def limit_layers(model: transformers.PreTrainedModel, n_layers: int) -> None:
     if hasattr(model, 'transformer'):
@@ -27,6 +449,7 @@ 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):
@@ -78,8 +501,7 @@ 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):
@@ -89,6 +511,9 @@ 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):
@@ -107,9 +532,11 @@ 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:
+
+        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 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, :]
@@ -137,12 +564,24 @@ 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, 
@@ -199,6 +638,7 @@ 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
@@ -233,7 +673,6 @@ class BiEncoder(transformers.PreTrainedModel):
             else:
                 document_embeddings = document_embeddings.max(dim=1)
         output = self.mlp(document_embeddings)
-        # breakpoint()
 
         if output_hidden_states:
             return {
@@ -258,9 +697,10 @@ class DatasetConditionedAutoregressive(transformers.PreTrainedModel, ContextualM
         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
+        self.input_ln = torch.nn.LayerNorm(
+            self.backbone_hidden_size, 
+            eps=1e-5
+        )
         
         # Override contextual init
         self.output_projection = torch.nn.Sequential(
@@ -286,7 +726,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?
-        print0("Warning: Positional embedding disabling not implemented for LLAMA.")
+        print("Warning: Positional embedding disabling not implemented for LLAMA.")
     
     def forward(
             self, 
@@ -312,6 +752,7 @@ 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(
@@ -333,34 +774,11 @@ class DatasetConditionedAutoregressive(transformers.PreTrainedModel, ContextualM
             output_hidden_states=True,
         ) # (1, 4 + b + s, d)
         # trim soft prompt
-        output_vectors = output.hidden_states[-1]
+        last_hidden_state = 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:]
+        output_vectors = last_hidden_state[:, n_soft_prompt_tokens:, :]
+        output_attention_mask = input_attention_mask[:, n_soft_prompt_tokens:]
 
         # Take last token position
         if vars(self.config).get("pooling_strategy") == "last_token":
@@ -371,6 +789,7 @@ 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:
@@ -382,6 +801,7 @@ class DatasetConditionedAutoregressive(transformers.PreTrainedModel, ContextualM
             return output
 
 
+
 class DatasetConditionedBiencoder(transformers.PreTrainedModel, ContextualModelMixin):
     def __init__(
             self, 
@@ -418,7 +838,7 @@ class DatasetConditionedBiencoder(transformers.PreTrainedModel, ContextualModelM
                 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}")
+            print0(f"modified {rotary_disabled} rotary modules – set rotary_start_pos to {rotary_start_pos}")
     
     def forward(
             self, 
@@ -547,7 +967,7 @@ class ContextualDocumentEmbeddingTransformer(transformers.PreTrainedModel):
         ):
         super().__init__(config=config)
         dataset_backbone, _ = load_embedder_and_tokenizer(
-            vars(config).get("dataset_backbone") or config.embedder
+            vars(config).get("dataset_backbone", config.embedder)
         )
 
         if config.limit_layers:
@@ -592,7 +1012,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(
@@ -606,17 +1026,12 @@ 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:
-        raise ValueError(f'unknown model cls {name}')
+        raise ValueError(f'unknown model cls {name}')