edit source

model.py

@@ -1,31 +1,7 @@
-###################################################################################################
-###################################################################################################
-###################################################################################################
-
-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 json
 import math
 import multiprocessing
 import os
@@ -34,7 +10,6 @@ import torch
 import torch.nn as nn
 import transformers
 
-
 class ContextualModelConfig(transformers.configuration_utils.PretrainedConfig):
     """We create a dummy configuration class that will just set properties
     based on whatever kwargs we pass in.
@@ -54,14 +29,13 @@ class ContextualModelConfig(transformers.configuration_utils.PretrainedConfig):
                 continue
         super().__init__()
 
-
 def load_embedder_and_tokenizer(name: str) -> Tuple[
         transformers.PreTrainedModel, 
         transformers.PreTrainedTokenizer
 ]:
-    print("Loading model:", name)
+    assert name is not None, "name must be provided to load_embedder_and_tokenizer"
     if name.startswith("nomic") or (name == "bert-base-uncased"):
-        model = ContextualNomicBertForPreTraining.from_pretrained(name, trust_remote_code=True).bert
+        model = transformers.AutoModelForMaskedLM.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(
@@ -106,8 +80,6 @@ def load_embedder_and_tokenizer(name: str) -> Tuple[
         #     from optimum.bettertransformer import BetterTransformer
         #     model = BetterTransformer.transform(model)
     return model, tokenizer
-
-
 def get_world_size() -> int:
     try:
         return torch.distributed.get_world_size()
@@ -318,7 +290,7 @@ def maxsim(
         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_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)
@@ -471,7 +443,6 @@ def disable_causality(model: torch.nn.Module):
         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:
@@ -511,9 +482,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 +500,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, :]
@@ -558,6 +524,8 @@ class ContextualModelMixin(nn.Module):
             null_embeddings = self.sequence_dropout_null_embedding[None, None].expand(batch_size, corpus_size, -1)
             dataset_embeddings = null_embeddings
         
+        # print(f"[ContextualModelMixin] dataset_embeddings.shape = {dataset_embeddings.shape}")
+        
         # 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)
@@ -630,8 +598,15 @@ class BiEncoder(transformers.PreTrainedModel):
                 [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 dataset_input_ids
+        # del dataset_attention_mask
         del token_type_ids
 
+        # from cde.lib.dist import get_rank
+        # tokenizer = transformers.AutoTokenizer.from_pretrained("bert-base-uncased")
+        # if get_rank() == 0:
+        #     breakpoint()
+        # torch.distributed.barrier()
         outputs = (
             self.embedder(
                 input_ids=input_ids,
@@ -801,7 +776,6 @@ class DatasetConditionedAutoregressive(transformers.PreTrainedModel, ContextualM
             return output
 
 
-
 class DatasetConditionedBiencoder(transformers.PreTrainedModel, ContextualModelMixin):
     def __init__(
             self, 
@@ -812,14 +786,12 @@ class DatasetConditionedBiencoder(transformers.PreTrainedModel, ContextualModelM
         self.backbone = dataset_backbone
         self.hidden_size = self.backbone.config.hidden_size
         self.hidden_size = dataset_backbone.config.hidden_size
+        # self.input_ln = torch.nn.LayerNorm(
+        #     self.hidden_size, 
+        #     eps=self.backbone.config.layer_norm_epsilon
+        # )
         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:
@@ -848,55 +820,48 @@ class DatasetConditionedBiencoder(transformers.PreTrainedModel, ContextualModelM
             output_hidden_states: bool = False,
             null_dataset_embedding: bool = False,
         ) -> torch.Tensor:
+        # print(f"[DatasetConditionedBiencoder - 0] input_ids.shape => {input_ids.shape} // dataset_embeddings.shape =", dataset_embeddings.shape)
         soft_prompt = self._prepare_dataset_embeddings(
             input_ids=input_ids,
             dataset_embeddings=dataset_embeddings,
             null_dataset_embedding=null_dataset_embedding,
         )
+        # print(f"[DatasetConditionedBiencoder - 1] soft_prompt.shape => {soft_prompt.shape}")
         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)
+        # print("[2] inputs_embeds.shape =", inputs_embeds.shape)
         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)
+        # print("[3.a] inputs_embeds.shape =", inputs_embeds.shape)
+        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,
+            attention_mask=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]
+        # print("n_soft_prompt_tokens =", 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)
-            )
-            output_attention_mask = torch.cat((backbone_attention_mask, attention_mask & ~is_instruction_token_mask), dim=1)
-        else:
-            output_attention_mask = input_attention_mask
+        output_vectors = output.last_hidden_state[:, n_soft_prompt_tokens:, :]
+        output_attention_mask = attention_mask[:, n_soft_prompt_tokens:]
 
-        if self.pool_ignore_contextual_tokens:
-            output_vectors = output_vectors[:, n_soft_prompt_tokens:, :]
-            output_attention_mask = output_attention_mask[:, n_soft_prompt_tokens:]
+        # print("pooling output_vectors.shape =", output_vectors.shape, "and output_attention_mask.shape =", output_attention_mask.shape)
         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
+        # TODO: Argparse for pooling strategy.
+        # output_vectors = torch.cat((soft_prompt_pooled, output_pooled), dim=1) # (b, d) + (b, d) -> (b, 2d)
+        # print("output_pooled.shape =", output_pooled.shape)
+        output = self.output_projection(output_pooled) # (b, 2d) -> (b, d)
+
+        # print("returning output.shape =", output.shape)
 
         if output_hidden_states:
             return {
@@ -967,7 +932,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:
@@ -1026,6 +991,8 @@ class ContextualDocumentEmbeddingTransformer(transformers.PreTrainedModel):
             output_hidden_states=output_hidden_states,
         )
 
+
+
 def get_model_class(name: str):
     if name in 'transductive': 
         return ContextualDocumentEmbeddingTransformer
@@ -1034,4 +1001,4 @@ def get_model_class(name: str):
     elif name == "dataset_prefix_biencoder":
         return DatasetPrefixBiencoder
     else:
-        raise ValueError(f'unknown model cls {name}')
+        raise ValueError(f'unknown model cls {name}')