quocviethere's picture
Update app.py
7a1c2c4 verified
import gradio as gr
from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
import torch
# -----------------------------
# Configuration Section
# -----------------------------
MODEL_NAME = "quocviethere/imdb-roberta" # Replace with your actual model ID
# -----------------------------
# Model Loading Section
# -----------------------------
try:
# Load tokenizer and model from Hugging Face Hub
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
# Initialize the sentiment analysis pipeline
sentiment_pipeline = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer)
# Verify label mapping
label_mapping = model.config.id2label
print(f"Model label mapping: {label_mapping}")
except Exception as e:
print(f"Error loading model: {e}")
raise
# -----------------------------
# Sentiment Analysis Function
# -----------------------------
def analyze_sentiment(text):
try:
# Perform sentiment analysis
result = sentiment_pipeline(text)[0]
# Extract label and score
label = result['label']
score = result['score']
# Map label to sentiment
if label in label_mapping.values():
sentiment = "Positive 😊" if label == "POSITIVE" else "Negative 😞"
else:
# Handle unexpected labels
sentiment = label
print(f"Unexpected label received: {label}")
confidence = f"Confidence: {round(score * 100, 2)}%"
return sentiment, confidence
except Exception as e:
print(f"Error during sentiment analysis: {e}")
return "Error", "Could not process the input."
# -----------------------------
# Gradio Interface Section
# -----------------------------
iface = gr.Interface(
fn=analyze_sentiment,
inputs=gr.Textbox(
lines=5,
placeholder="Enter a movie review here...",
label="Movie Review"
),
outputs=[
gr.Textbox(label="Sentiment"),
gr.Textbox(label="Confidence")
],
title="IMDb Sentiment Analysis with RoBERTa",
description="Analyze the sentiment of movie reviews using a fine-tuned RoBERTa model.",
examples=[
["I loved the cinematography and the story was captivating."],
["The movie was a complete waste of time. Poor acting and boring plot."]
],
theme="default"
)
# -----------------------------
# Launch the Interface
# -----------------------------
iface.launch()