Upload folder using huggingface_hub

config.json

@@ -23,5 +23,8 @@
   "transformers_version": "4.51.3",
   "type_vocab_size": 2,
   "use_cache": true,
-  "vocab_size": 32768
-}
+  "vocab_size": 32768,
+  "auto_map": {
+    "AutoModel": "modeling_simcse.SimCSEInferenceModel"
+  }
+}

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d17f22fdf6834dd2fae0baad70d8a15bc36644056ba215f582a5a4b5c4012b4c
+oid sha256:c89b84d9d78f111b2c00bb3bda2063b07ffa943fc59581b095c6cab6fd4b181b
 size 894432952

modeling_simcse.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+from transformers import BertModel, BertConfig, PreTrainedModel
+from transformers.tokenization_utils_base import BatchEncoding
+import torch, torch.nn as nn, torch.nn.functional as F
+
+class SimCSEInferenceModel(PreTrainedModel):
+    config_class = BertConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        base_cfg = BertConfig(**config.to_dict())
+        self.encoder_input  = BertModel(base_cfg)
+        self.encoder_output = BertModel(base_cfg)
+
+        hidden = self.encoder_input.config.hidden_size
+        self.dense_input  = nn.Linear(hidden, hidden)
+        self.dense_output = nn.Linear(hidden, hidden)
+        self.activation   = nn.Tanh()
+        self.temperature  = getattr(config, "simcse_temperature", 0.05)
+
+    @torch.no_grad()
+    def encode_input(self, tok: BatchEncoding) -> torch.Tensor:
+        h = self.encoder_input(**tok).last_hidden_state[:, 0]
+        return self.activation(self.dense_input(h))
+
+    @torch.no_grad()
+    def encode_output(self, tok: BatchEncoding) -> torch.Tensor:
+        h = self.encoder_output(**tok).last_hidden_state[:, 0]
+        return self.activation(self.dense_output(h))
+
+    def forward(self, tokenized_texts_1: BatchEncoding, tokenized_texts_2: BatchEncoding, labels: torch.Tensor, **_):
+        device = next(self.parameters()).device
+        z1 = F.normalize(self.encode_input(tokenized_texts_1.to(device)), dim=-1)
+        z2 = F.normalize(self.encode_output(tokenized_texts_2.to(device)), dim=-1)
+        sim = torch.matmul(z1, z2.T)
+        loss = F.cross_entropy(sim / self.temperature, labels.to(device))
+        return {"loss": loss, "logits": sim}