File size: 2,535 Bytes
c8375ec 3c49916 |
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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
---
license: mit
datasets:
- KomeijiForce/Inbedder-Pretrain-Data
language:
- en
---
```python
import torch
from torch import nn
from torch.nn.functional import gelu, cosine_similarity
from transformers import AutoTokenizer, AutoModel, AutoModelForMaskedLM
import numpy as np
class InBedder():
def __init__(self, path='KomeijiForce/inbedder-roberta-large', device='cuda:0'):
model = AutoModelForMaskedLM.from_pretrained(path)
self.tokenizer = AutoTokenizer.from_pretrained(path)
self.model = model.roberta
self.dense = model.lm_head.dense
self.layer_norm = model.lm_head.layer_norm
self.device = torch.device(device)
self.model = self.model.to(self.device)
self.dense = self.dense.to(self.device)
self.layer_norm = self.layer_norm.to(self.device)
self.vocab = self.tokenizer.get_vocab()
self.vocab = {self.vocab[key]:key for key in self.vocab}
def encode(self, input_texts, instruction, n_mask):
if type(instruction) == str:
prompts = [instruction + self.tokenizer.mask_token*n_mask for input_text in input_texts]
elif type(instruction) == list:
prompts = [inst + self.tokenizer.mask_token*n_mask for inst in instruction]
inputs = self.tokenizer(input_texts, prompts, padding=True, truncation=True, return_tensors='pt').to(self.device)
mask = inputs.input_ids.eq(self.tokenizer.mask_token_id)
outputs = self.model(**inputs)
logits = outputs.last_hidden_state[mask]
logits = self.layer_norm(gelu(self.dense(logits)))
logits = logits.reshape(len(input_texts), n_mask, -1)
logits = logits.mean(1)
logits = (logits - logits.mean(1, keepdim=True)) / logits.std(1, keepdim=True)
return logits
inbedder = InBedder(path='KomeijiForce/inbedder-roberta-large', device='cpu')
texts = ["I love cat!", "I love dog!", "I dislike cat!"]
instruction = "What is the animal mentioned here?"
embeddings = inbedder.encode(texts, instruction, 3)
cosine_similarity(embeddings[:1], embeddings[1:], dim=1)
# tensor([0.9374, 0.9917], grad_fn=<SumBackward1>)
texts = ["I love cat!", "I love dog!", "I dislike cat!"]
instruction = "What is emotion expressed here?"
embeddings = inbedder.encode(texts, instruction, 3)
cosine_similarity(embeddings[:1], embeddings[1:], dim=1)
# tensor([0.9859, 0.8537], grad_fn=<SumBackward1>)
``` |