File size: 1,454 Bytes
6fcf1a4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
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
|