Spaces:
Running
Running
File size: 1,928 Bytes
f584171 2983032 1f9fb90 2983032 d98a217 f584171 dc9c50d f584171 d98a217 f584171 d98a217 f584171 dc9c50d f584171 dc9c50d 17962cc f584171 70d5d1d 1f9fb90 de1e086 1f9fb90 d98a217 f584171 2983032 f584171 d98a217 2983032 1f9fb90 2983032 d98a217 f584171 5a35002 |
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 |
import torch
from transformers import BertForSequenceClassification, AutoTokenizer
from extract import summarize_text as predict_summarization
# path_emo = 'Djacon/rubert-tiny2-russian-emotion-detection'
path_emo = './models/emotion_detection/'
model_emo = BertForSequenceClassification.from_pretrained(path_emo)
tokenizer_emo = AutoTokenizer.from_pretrained(path_emo)
LABELS = ['Neutral', 'Joy', 'Sadness', 'Anger', 'Enthusiasm', 'Surprise', 'Disgust', 'Fear', 'Guilt', 'Shame']
# Probabilistic prediction of emotion in a text
@torch.no_grad()
def predict_emotions(text: str) -> str:
inputs = tokenizer_emo(text, max_length=512, truncation=True,
return_tensors='pt')
outputs = model_emo(**inputs)
pred = torch.nn.functional.sigmoid(outputs.logits)
emotions_list = {}
for i in range(len(pred[0].tolist())):
emotions_list[LABELS[i]] = round(100 * pred[0].tolist()[i], 2)
return '\n'.join(f"{k}: {v}%" for k, v in sorted(emotions_list.items(),
key=lambda x: -x[1]))
# path_gram = 'Djacon/mbert-gram'
path_gram = './models/mbert-gram/'
model_gram = BertForSequenceClassification.from_pretrained(path_gram)
tokenizer_gram = AutoTokenizer.from_pretrained(path_gram)
@torch.no_grad()
def predict_acceptance(text: str) -> str:
inputs = tokenizer_gram(text, truncation=True, return_tensors='pt')
output = model_gram(**inputs).logits[0]
pred = torch.nn.functional.softmax(output, dim=0)[1]
return f'Acceptance: {100 * pred.item():.2f}%'
def test():
predict_emotions('I am so happy now!')
print('\n>>> Emotion Detection successfully initialized! <<<')
predict_summarization('I am so happy now!')
print('>>> Extractive Summarizer successfully initialized! <<<')
predict_acceptance('I am so happy now!')
print('>>> Grammar Checker successfully initialized! <<<\n')
test()
|