shae2977's picture
add Hugging Face login
585839c
Raw
History Blame Contribute Delete
2.66 kB
from transformers import XLMRobertaTokenizer, XLMRobertaForSequenceClassification
from vaderSentiment.vaderSentiment import SentimentIntensityAnalyzer
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from huggingface_hub import login
import logging
import torch
import os
login(token=os.getenv("HF_TOKEN"))
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(levelname)s - %(message)s",
handlers=
[logging.FileHandler('app.log'),
logging.StreamHandler()]
)
logger = logging.getLogger(__name__)
model = XLMRobertaForSequenceClassification.from_pretrained('shae2977/xlm-roberta-hinglish-sentiment-analysis')
tokenizer = XLMRobertaTokenizer.from_pretrained('shae2977/xlm-roberta-hinglish-sentiment-analysis')
vader = SentimentIntensityAnalyzer()
def return_sentiment(comment):
score = vader.polarity_scores(comment)['compound']
label = None
if score >= 0.05:
label = 'Positive'
elif score <= -0.05:
label = 'Negative'
else:
label = 'Neutral'
return label
app = FastAPI()
class commentsinput(BaseModel):
comments : list[str]
@app.get('/')
def home():
return {'message' : 'Server is running!'}
@app.post('/predict')
def predict(input: commentsinput):
if not input.comments:
logger.error('Comment field empty')
raise HTTPException(status_code=400, detail='Comments field empty')
logger.info("Comment batch of %s comments recieved",len(input.comments))
try:
#roberta
results = []
batch_size = 50
for i in range(0, len(input.comments),batch_size):
batch = input.comments[i:batch_size+i]
comment_tokens = tokenizer(batch, padding=True, truncation=True, return_tensors='pt')
outputs = model(**comment_tokens)
pred = torch.argmax(outputs.logits, dim=1)
predictions = pred.numpy().tolist()
results.extend(predictions)
label_map = {0: 'Negative', 1: 'Neutral', 2: 'Positive'}
roberta_preds = []
for r in results:
roberta_preds.append(label_map[r])
#vader
vader_preds = []
for comment in input.comments:
vader_preds.append(return_sentiment(comment))
logger.info('Prediction successful')
return {'xlm-ROBERTa Predictions' : roberta_preds,
'VADER Predictions' : vader_preds}
except Exception as e:
logger.error('Prediction failed: %s',e)
raise HTTPException(status_code=500, detail=e)