karan99300's picture
Update app.py
8d76671 verified
import torch
from transformers import BertTokenizer, BertForSequenceClassification
import gradio as gr
label_dict={'neutral': 0,'negative': 1, 'positive': 2}
model = BertForSequenceClassification.from_pretrained("bert-base-uncased",
num_labels=len(label_dict),
output_attentions=False,
output_hidden_states=False)
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased',
do_lower_case=True)
model.load_state_dict(torch.load('finetuned_BERT_epoch_2.model',map_location='cpu'))
model.eval()
def get_key_by_value(dictionary, target_value):
for key, value in dictionary.items():
if value == target_value:
return key
def predict_sentiment(text):
inputs = tokenizer(text, return_tensors="pt")
inputs.to('cpu')
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
probabilities = torch.nn.functional.softmax(logits, dim=1)
predicted_class = torch.argmax(probabilities, dim=1).item()
return get_key_by_value(label_dict,predicted_class)
iface = gr.Interface(
fn=predict_sentiment,
inputs=gr.Textbox(),
outputs=gr.Textbox(),
live=True,
title="BERT Sentiment Analysis (CPU)",
description="Enter a text and get sentiment prediction.",
)
iface.launch()