import gradio as gr
import torch
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification

# Load the model and tokenizer
model_name = "AventIQ-AI/distilbert-base-uncased-sentiment-analysis"
tokenizer = DistilBertTokenizer.from_pretrained(model_name)
model = DistilBertForSequenceClassification.from_pretrained(model_name)

def predict_sentiment(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    with torch.no_grad():
        logits = model(**inputs).logits
    predicted_class_id = torch.argmax(logits, dim=-1).item()
    sentiment = "Positive" if predicted_class_id == 1 else "Negative"
    return sentiment

# Create Gradio interface
iface = gr.Interface(
    fn=predict_sentiment,
    inputs=gr.Textbox(lines=3, placeholder="Enter text for sentiment analysis..."),
    outputs=gr.Textbox(label="Sentiment"),
    title="DistilBERT Sentiment Analysis",
    description="Enter a sentence to classify its sentiment as Positive or Negative using a fine-tuned DistilBERT model.",
)

if __name__ == "__main__":
    iface.launch()