NLP / sentiment.py
ashishraics's picture
updated sentiment classifier
b72368a
raw
history blame
No virus
712 Bytes
import torch
from transformers import AutoModelForSequenceClassification,AutoTokenizer
model=AutoModelForSequenceClassification.from_pretrained('sentiment_classifier/')
tokenizer=AutoTokenizer.from_pretrained('sentiment_classifier/')
def classify_sentiment(texts,model=model,tokenizer=tokenizer):
"""
user will pass texts separated by comma
"""
try:
texts=texts.split(',')
except:
pass
input = tokenizer(texts, padding=True, truncation=True,
return_tensors="pt")
logits = model(**input)['logits'].softmax(dim=1)
logits = torch.argmax(logits, dim=1)
output = ['Positive' if i == 1 else 'Negative' for i in logits]
return output