Initial commit

cust_transformer.py

@@ -0,0 +1,47 @@
+
+from sentence_transformers import models
+import torch
+import torch.nn as nn
+
+class CustTrans(models.Transformer):
+
+  def __init__(self, *args, **kwargs):
+      super().__init__(*args, **kwargs)
+      self.curr_task_type = None
+      self._rebuild_taskembedding(['sts', 'quora'])
+
+  def forward(self, inputs, task_type=None):
+
+    enc = self.auto_model(**inputs).last_hidden_state
+
+    if task_type == None:
+      task_type = self.curr_task_type
+
+    if  task_type in self.task_types:
+      idx = torch.tensor(self.task_types.index(task_type), device=self.TaskEmbedding.weight.device)
+      hyp = self.TaskEmbedding(idx)
+      inputs['token_embeddings'] = self._project(enc, hyp)
+
+    else:
+      inputs['token_embeddings'] = enc
+
+    return inputs
+
+  def _set_curr_task_type(self, task_type):
+    self.curr_task_type = task_type
+
+  def _set_taskembedding_grad(self, value):
+        self.TaskEmbedding.weight.requires_grad = value
+
+  def _set_transformer_grad(self, value):
+        for param in self.auto_model.parameters():
+          param.requires_grad = value
+
+  def _rebuild_taskembedding(self, task_types):
+    self.task_types = task_types
+    self.task_emb = 1 - torch.eye(len(self.task_types),768)
+    self.TaskEmbedding = nn.Embedding(len(self.task_types), 768).from_pretrained(self.task_emb)
+
+  def _project(self, v, normal_hyper):
+    # return v - torch.dot(v, normal_hyper)*normal_hyper / torch.norm(normal_hyper)**2
+    return v*normal_hyper

modules.json

@@ -3,7 +3,7 @@
     "idx": 0,
     "name": "0",
     "path": "",
-    "type": "custom_trans.CustTrans"
+    "type": "cust_transformer.CustTrans"
   },
   {
     "idx": 1,