Spaces:
Running on Zero
Running on Zero
File size: 2,050 Bytes
b964596 | 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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | from transformers import (
AutoTokenizer,
AutoModelForSequenceClassification
)
import torch
MODEL_NAME = "jvomiranda/BERTimbau-Sent-Analysis"
class BertimbauModel:
def __init__(self):
self.device = torch.device(
"cuda"
if torch.cuda.is_available()
else "cpu"
)
self.tokenizer = AutoTokenizer.from_pretrained(
MODEL_NAME
)
self.model = (
AutoModelForSequenceClassification
.from_pretrained(MODEL_NAME)
.to(self.device)
)
self.model.eval()
@torch.no_grad()
def predict(self, text: str):
encoding = self.tokenizer(
text,
return_tensors="pt",
truncation=True,
max_length=128
)
encoding = {
key: value.to(self.device)
for key, value in encoding.items()
}
outputs = self.model(**encoding, output_hidden_states=True)
probabilities = torch.softmax(
outputs.logits,
dim=1
)[0]
prediction = torch.argmax(
probabilities
).item()
tokens = self.tokenizer.convert_ids_to_tokens(
encoding["input_ids"][0]
)
token_ids = (
encoding["input_ids"][0]
.cpu()
.tolist()
)
embeddings = outputs.hidden_states[-1][0]
representation = []
for token, vector in zip(tokens, embeddings):
representation.append({
"token": token,
"vector": [
round(float(x), 2)
for x in vector[:5]
]
})
return {
"prediction": prediction,
"probabilities": {
0: float(probabilities[0]),
1: float(probabilities[1])
},
"tokens": tokens,
"token_ids": token_ids,
"representation": representation
} |