FPXDoinb's picture
Update app.py
36b8515 verified
import torch
from transformers import BertTokenizer
from torch.nn.functional import softmax
from google.colab import drive
import gradio as gr
drive.mount('/content/drive')
# Set the correct path for the model within the Hugging Face Space
model = torch.load('/content/drive/My Drive/Emotion/emotion_model.pth')
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model.eval() # Set the model to evaluation mode
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
def predict_emotions(text):
inputs = tokenizer.encode_plus(text, return_tensors="pt", max_length=512, truncation=True, padding='max_length')
input_ids = inputs['input_ids'].to(device)
attention_mask = inputs['attention_mask'].to(device)
with torch.no_grad():
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
probabilities = softmax(outputs.logits, dim=-1).squeeze()
emotions = ['Sadness', 'Joy', 'Love', 'Anger', 'Fear', 'Surprise']
response = ", ".join(f"{emotion}: {prob * 100:.2f}%" for emotion, prob in zip(emotions, probabilities))
return response
iface = gr.Interface(fn=predict_emotions, inputs="text", outputs="label")
iface.launch()