adding sentence-transformers specific files

0_Model/config.json

@@ -0,0 +1,3 @@
+{
+    "model_name_or_path": "/mnt/mydrive/embedding-models-fine-tuned/merged_models/7B_EU24_2_5t_with_BIG_dataset_multiple_neg_full_model"
+}

1_TokenPooling/config.json

@@ -0,0 +1,3 @@
+{
+    "dimension": 4096
+}

config_sentence_transformers.json

@@ -0,0 +1,10 @@
+{
+  "__version__": {
+    "sentence_transformers": "3.2.1",
+    "transformers": "4.41.2",
+    "pytorch": "2.4.0+cu121"
+  },
+  "prompts": {},
+  "default_prompt_name": null,
+  "similarity_fn_name": null
+}

embedding_model.py

@@ -0,0 +1,77 @@
+import json
+import os
+from typing import Union, List, Dict, Tuple
+
+import torch
+from sentence_transformers import models
+
+MODEL_PATH_IN_MODULE = "0_Model"
+
+
+class EmbeddingModel(models.Transformer):
+    def __init__(self, model_name_or_path: str, *args, **kwargs):
+        self.model_name_or_path = model_name_or_path
+        super().__init__(model_name_or_path, *args, **kwargs)
+
+    def tokenize(
+        self,
+        texts: Union[List[str], List[Dict], List[Tuple[str, str]]],
+        padding: Union[str, bool] = True,
+    ) -> Dict[str, torch.Tensor]:
+        """Tokenizes a text and maps tokens to token-ids"""
+        output = {}
+        if isinstance(texts[0], str):
+            texts = [x + self.tokenizer.eos_token for x in texts]
+            to_tokenize = [texts]
+        elif isinstance(texts[0], dict):
+            to_tokenize = []
+            output["text_keys"] = []
+            for lookup in texts:
+                text_key, text = next(iter(lookup.items()))
+                to_tokenize.append(text)
+                output["text_keys"].append(text_key)
+            to_tokenize = [to_tokenize]
+        else:
+            batch1, batch2 = [], []
+            for text_tuple in texts:
+                batch1.append(text_tuple[0])
+                batch2.append(text_tuple[1])
+            to_tokenize = [batch1, batch2]
+
+        output.update(
+            self.tokenizer(
+                *to_tokenize,
+                padding="max_length",
+                truncation=True,
+                return_tensors="pt",
+                max_length=512,
+            )
+        )
+
+        # this is specific to OpenGPT-X model
+        output.pop("token_type_ids", None)
+
+        return output
+
+    def get_config_dict(self) -> dict[str, str]:
+        return {"model_name_or_path": self.model_name_or_path}
+
+    def save(self, save_dir: str, **kwargs) -> None:
+        # self.auto_model.save_pretrained(save_dir, safe_serialization=True)
+        # self.tokenizer.save_pretrained(save_dir)
+
+        with open(os.path.join(save_dir, "sentence_bert_config.json"), "w+") as f:
+            json.dump(self.get_config_dict(), f, indent=4)
+
+        model_path = os.path.join(save_dir, MODEL_PATH_IN_MODULE)
+        if not os.path.exists(model_path):
+            os.makedirs(model_path)
+
+        with open(f"{model_path}/config.json", "w+") as f:
+            json.dump(self.get_config_dict(), f, indent=4)
+
+    @staticmethod
+    def load(load_dir: str, **kwargs) -> "EmbeddingModel":
+        with open(os.path.join(load_dir, "config.json")) as fIn:
+            config = json.load(fIn)
+        return EmbeddingModel(**config)

modules.json

@@ -0,0 +1,20 @@
+[
+    {
+        "idx": 0,
+        "name": "0",
+        "path": "0_Model",
+        "type": "embedding_model.EmbeddingModel"
+    },
+    {
+        "idx": 1,
+        "name": "1",
+        "path": "1_TokenPooling",
+        "type": "token_pooling.TokenPooling"
+    },
+    {
+        "idx": 2,
+        "name": "2",
+        "path": "2_Normalize",
+        "type": "sentence_transformers.models.Normalize"
+    }
+]

sentence_bert_config.json

@@ -0,0 +1,3 @@
+{
+    "model_name_or_path": "/mnt/mydrive/embedding-models-fine-tuned/merged_models/7B_EU24_2_5t_with_BIG_dataset_multiple_neg_full_model"
+}

token_pooling.py

@@ -0,0 +1,59 @@
+import json
+import os
+import torch
+import torch.nn as nn
+
+
+class TokenPooling(nn.Module):
+    def __init__(self, dimension: int = 4096) -> None:
+        super(TokenPooling, self).__init__()
+        self.dimension = dimension
+
+    def forward(
+        self, features: dict[str, torch.Tensor], **kwargs
+    ) -> dict[str, torch.Tensor]:
+        token_embeddings = features["token_embeddings"]
+        attention_mask = features["attention_mask"]
+
+        embeddings = self.pool(
+            last_hidden_state=token_embeddings, attention_mask=attention_mask
+        )
+        features["sentence_embedding"] = embeddings
+        return features
+
+    def pool(
+        self, last_hidden_state: torch.Tensor, attention_mask: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Here, we take the embedding of the last token from the last layer
+        """
+        left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
+        if left_padding:
+            return last_hidden_state[:, -1]
+        else:
+            sequence_lengths = attention_mask.sum(dim=1) - 1
+            batch_size = last_hidden_state.shape[0]
+            return last_hidden_state[
+                torch.arange(batch_size, device=last_hidden_state.device).long(),
+                sequence_lengths.long(),
+            ]
+
+    def get_sentence_embedding_dimension(self) -> int:
+        return self.dimension
+
+    def get_config_dict(self) -> dict[str, float]:
+        return {"dimension": self.dimension}
+
+    def save(self, save_dir: str, **kwargs) -> None:
+        pooling_path = os.path.join(save_dir)
+        if not os.path.exists(pooling_path):
+            os.makedirs(pooling_path)
+
+        with open(f"{pooling_path}/config.json", "w+") as f:
+            json.dump(self.get_config_dict(), f, indent=4)
+
+    @staticmethod
+    def load(load_dir: str, **kwargs) -> "TokenPooling":
+        with open(os.path.join(load_dir, "config.json")) as fIn:
+            config = json.load(fIn)
+        return TokenPooling(**config)