Upload ZeroShotEmbedding

config.json

@@ -0,0 +1,16 @@
+{
+  "architectures": [
+    "ZeroShotEmbedding"
+  ],
+  "auto_map": {
+    "AutoConfig": "model.ZeroShotEmbeddingConfig",
+    "AutoModel": "model.ZeroShotEmbedding"
+  },
+  "base_embedding_model": "all-mpnet-base-v2",
+  "hidden_size": 2048,
+  "input_size": 768,
+  "model_type": "embedding-head",
+  "output_size": 128,
+  "torch_dtype": "float32",
+  "transformers_version": "4.35.0"
+}

model.py

@@ -0,0 +1,90 @@
+from transformers import PreTrainedModel, PretrainedConfig
+from sentence_transformers import SentenceTransformer
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+class ZeroShotEmbeddingConfig(PretrainedConfig):
+    model_type = "embedding-head"
+
+    def __init__(self, input_size=768, hidden_size=2048, output_size=128, base_embedding_model='all-mpnet-base-v2', **kwargs):
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.output_size = output_size
+        self.base_embedding_model = base_embedding_model
+        super().__init__(**kwargs)
+
+
+class ZeroShotEmbedding(PreTrainedModel):
+    def __init__(self, config):
+        super(ZeroShotEmbedding, self).__init__(config)
+
+        input_size = config.input_size
+        hidden_size = config.hidden_size
+        output_size = config.output_size
+
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.output_size = output_size
+        # 3-layer MLP: input embedding -> hidden -> output embedding
+        self.fc1 = nn.Linear(input_size * 2, hidden_size)
+        self.fc2 = nn.Linear(hidden_size, output_size)
+        self.gelu = nn.GELU()
+
+    def forward(self, prompt_embedding, text_a_embedding, text_b_embedding=None, labels=None, **kwargs):
+        # document_embedding: [batch_size, input_size]
+        # prompt_embedding: [batch_size, input_size]
+        # output: [batch_size, output_size]
+
+        # concatenate document embedding and prompt embedding
+        # [batch_size, input_size * 2]
+        x = torch.cat((text_a_embedding, prompt_embedding), dim=1)
+        if text_b_embedding is not None:
+            # concatenate document embedding and prompt embedding
+            # [batch_size, input_size * 2]
+            x2 = torch.cat((text_b_embedding, prompt_embedding), dim=1)
+
+        # 3-layer MLP
+        x = self.fc1(x)
+        x = self.gelu(x)
+        x = self.fc2(x)
+        x = nn.functional.normalize(x, p=2, dim=1)
+        if text_b_embedding is not None:
+            x2 = self.fc1(x2)
+            x2 = self.gelu(x2)
+            x2 = self.fc2(x2)
+            x2 = nn.functional.normalize(x2, p=2, dim=1)
+            # Compute dot product for batches of output vectors
+            dot_product = torch.bmm(x.unsqueeze(1), x2.unsqueeze(2)).squeeze()
+            if labels is not None:
+                # Compute loss (magnitude of dot product minus label)
+                loss = torch.mean((dot_product - labels) ** 2)
+                return loss, dot_product
+            return dot_product
+        return x
+
+
+class ZeroShotEmbeddingForClustering(PreTrainedModel):
+    def __init__(self, config):
+        super(ZeroShotEmbeddingForClustering, self).__init__(config)
+        self.base_embedding_model = SentenceTransformer(
+            config.base_embedding_model)
+        self.head_model = ZeroShotEmbedding(config)
+
+    def forward(self, texts, prompt, **kwargs):
+        text_embeddings = self.base_embedding_model.encode(texts)
+        prompt_embedding = self.base_embedding_model.encode(prompt)
+        prompt_embeddings = np.tile(prompt_embedding, (len(texts), 1))
+        text_embeddings = torch.tensor(text_embeddings)
+        prompt_embeddings = torch.tensor(prompt_embeddings)
+        prompted_embeddings = self.head_model(
+            prompt_embeddings, text_embeddings)
+        similarity = torch.mm(prompted_embeddings,
+                              prompted_embeddings.transpose(0, 1))
+        return similarity
+
+
+ZeroShotEmbeddingConfig.register_for_auto_class()
+ZeroShotEmbedding.register_for_auto_class("AutoModel")
+ZeroShotEmbeddingForClustering.register_for_auto_class("AutoModel")

model.safetensors

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