import gradio as gr
from transformers import pipeline, BertTokenizerFast, BertForSequenceClassification

model_id = "AtulDeshpande/reddit_sentiment_model"

label_map = {"LABEL_0":"Negative", "LABEL_1":"Neutral", "LABEL_2": "Positive"}

tokenizer = BertTokenizerFast.from_pretrained(model_id)
model = BertForSequenceClassification.from_pretrained(model_id)
sentiment_pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)

def predict(text):
    result = sentiment_pipeline(text)[0]
    label = result['label']
    label = label_map[label]
    return f"{label} (confidence: {result['score']:.2f})"

demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=3, placeholder="Enter a Reddit comment here..."),
    outputs="text",
    title="Reddit Sentiment Classifier",
    description="Fine-tuned BERT model to classify Reddit comments as Positive, Neutral, or Negative."
)

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