Upload model

config.json

@@ -2,6 +2,10 @@
   "architectures": [
     "SiamseNNModel"
   ],
+  "auto_map": {
+    "AutoConfig": "configuration_siamese.SiameseConfig",
+    "AutoModel": "modeling_siamese.SiamseNNModel"
+  },
   "model_type": "AutoModel",
   "torch_dtype": "float32",
   "transformers_version": "4.24.0"

configuration_siamese.py

@@ -0,0 +1,14 @@
+from transformers import PretrainedConfig
+
+
+class SiameseConfig(PretrainedConfig):
+    model_type = "AutoModel"
+
+    def __init__(
+        self,
+        **kwargs):
+        super().__init__()
+        
+
+siamese_config = SiameseConfig()
+siamese_config.save_pretrained('siamse_nn')

modeling_siamese.py

@@ -0,0 +1,50 @@
+from transformers import PreTrainedModel, BertModel
+import torch
+
+checkpoint = 'cointegrated/rubert-tiny'
+
+class Lambda(torch.nn.Module):
+    def __init__(self, lambd):
+        super().__init__()
+        self.lambd = lambd
+    
+    def forward(self, x):
+         return self.lambd(x)
+
+
+class SiameseNN(torch.nn.Module):
+    def __init__(self):
+        super(SiameseNN, self).__init__()
+        l1_norm = lambda x: 1 - torch.abs(x[0] - x[1])
+        self.encoder = BertModel.from_pretrained(checkpoint)
+        self.merged = Lambda(l1_norm)
+        self.fc1 = torch.nn.Linear(312, 2)
+        self.softmax = torch.nn.Softmax()
+
+    
+    def forward(self, x):
+        first_encoded = self.encoder(**x[0]).pooler_output
+        #print("First: ", first_encoded)
+        second_encoded = self.encoder(**x[1]).pooler_output
+        l1_distance = self.merged([first_encoded, second_encoded])
+        #print(l1_distance.shape)
+        fc1 = self.fc1(l1_distance)
+        fc1 = self.softmax(fc1)
+        return fc1
+
+second_model = SiameseNN()
+second_model.load_state_dict(torch.load('siamese_state'))
+
+class SiamseNNModel(PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = second_model
+
+    
+    def forward(self, tensor, labels=None):
+        logits = self.model(tensor)
+        if labels is not None:
+            loss_fn = torch.nn.CrossEntropyLoss()
+            loss = loss_fn(logits, labels)
+            return {'loss': loss, 'logits': logits}
+        return {'logits': logits}