File size: 2,724 Bytes
8578816 |
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 |
from ast import literal_eval
import torch
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForTokenClassification, AutoModelForSequenceClassification
from transformers import BertForSequenceClassification, BertTokenizer, BertConfig
from math import exp
from . import label
class Model(object):
def __init__(self) -> None:
self.model_name = "indolem/indobert-base-uncased"
self.tokenizer = None
self.model = None
self.config = None
def load_model(self, model_name: str = None, tasks: str = None):
print(model_name)
if tasks == "emotion":
self.config = BertConfig.from_pretrained(model_name)
self.tokenizer = BertTokenizer.from_pretrained(model_name) \
if tasks == "emotion" else \
AutoTokenizer.from_pretrained(model_name)
if tasks == "emotion":
self.model = BertForSequenceClassification.from_pretrained(model_name, config=self.config)
elif tasks == "ner":
self.model = AutoModelForTokenClassification.from_pretrained(model_name)
else:
self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
def predict(self, sentences, tasks: str = None):
encoded_input = self.tokenizer(sentences,
return_tensors="pt",
padding=True,
truncation=True)
with torch.no_grad():
if tasks in ["emotion", "sentiment"]:
outputs = self.model(**encoded_input)
predicted_class = torch.argmax(outputs.logits, dim=1).item()
logits = outputs.logits.numpy()
probability = [exp(output)/(1+exp(output)) for output in logits[0]]
else:
recognizer = pipeline("token-classification", model=self.model, tokenizer=self.tokenizer)
outputs = recognizer(sentences)
if tasks in ["emotion", "sentiment"]:
result = {"label": label[tasks][predicted_class],
"score": probability[predicted_class]}
elif tasks == "ner":
result = []
for output in outputs:
result.append(
{
"entity": output["entity"],
"score": float(output["score"]),
"index": int(output["index"]),
"word": output["word"],
"start": int(output["start"]),
"end": int(output["end"])
}
)
else:
result = ""
return result |