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()