roycett's picture
Update app.py
5784c8d verified
import gradio as gr
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
model = AutoModelForSequenceClassification.from_pretrained("roycett/fine-tuned-roberta-yt-sentiment")
tokenizer = AutoTokenizer.from_pretrained("roycett/fine-tuned-roberta-yt-sentiment")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
def predict_sentiment(transcript):
inputs = tokenizer(transcript, return_tensors="pt", max_length=100, truncation=True, padding="max_length")
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
probabilities = torch.softmax(logits, dim=1)
prediction = torch.argmax(probabilities, dim=1)
label = "positive" if prediction.item() == 1 else "negative"
probability = probabilities[0][prediction.item()].item()
return f"Sentiment: {label} (Confidence: {probability:.2f})"
iface = gr.Interface(
fn=predict_sentiment,
inputs=gr.Textbox(lines=5, placeholder="Enter your transcript here..."),
outputs="text",
title="YouTube Transcript Sentiment Analysis",
description="Paste a YouTube transcript to analyze its sentiment."
)
iface.launch()