Views
No views yet
1import torch
2from torch import nn
3from torch.nn.functional import gelu, cosine_similarity
4from transformers import AutoTokenizer, AutoModel, AutoModelForMaskedLM
5
6import numpy as np
7
8class InBedder():
9
10 def __init__(self, path='KomeijiForce/inbedder-roberta-large', device='cuda:0'):
11
12 model = AutoModelForMaskedLM.from_pretrained(path)
13
14 self.tokenizer = AutoTokenizer.from_pretrained(path)
15 self.model = model.roberta
16 self.dense = model.lm_head.dense
17 self.layer_norm = model.lm_head.layer_norm
18
19 self.device = torch.device(device)
20 self.model = self.model.to(self.device)
21 self.dense = self.dense.to(self.device)
22 self.layer_norm = self.layer_norm.to(self.device)
23
24 self.vocab = self.tokenizer.get_vocab()
25 self.vocab = {self.vocab[key]:key for key in self.vocab}
26
27 def encode(self, input_texts, instruction, n_mask):
28
29 if type(instruction) == str:
30 prompts = [instruction + self.tokenizer.mask_token*n_mask for input_text in input_texts]
31 elif type(instruction) == list:
32 prompts = [inst + self.tokenizer.mask_token*n_mask for inst in instruction]
33
34 inputs = self.tokenizer(input_texts, prompts, padding=True, truncation=True, return_tensors='pt').to(self.device)
35
36 mask = inputs.input_ids.eq(self.tokenizer.mask_token_id)
37
38 outputs = self.model(**inputs)
39
40 logits = outputs.last_hidden_state[mask]
41
42 logits = self.layer_norm(gelu(self.dense(logits)))
43
44 logits = logits.reshape(len(input_texts), n_mask, -1)
45
46 logits = logits.mean(1)
47
48 logits = (logits - logits.mean(1, keepdim=True)) / logits.std(1, keepdim=True)
49
50 return logits
51
52inbedder = InBedder(path='KomeijiForce/inbedder-roberta-large', device='cpu')
53
54texts = ["I love cat!", "I love dog!", "I dislike cat!"]
55instruction = "What is the animal mentioned here?"
56embeddings = inbedder.encode(texts, instruction, 3)
57
58cosine_similarity(embeddings[:1], embeddings[1:], dim=1)
59# tensor([0.9374, 0.9917], grad_fn=<SumBackward1>)
60
61texts = ["I love cat!", "I love dog!", "I dislike cat!"]
62instruction = "What is emotion expressed here?"
63embeddings = inbedder.encode(texts, instruction, 3)
64
65cosine_similarity(embeddings[:1], embeddings[1:], dim=1)
66# tensor([0.9859, 0.8537], grad_fn=<SumBackward1>)