import gradio as gr
import numpy as np
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig
from scipy.special import softmax

# Load fine-tuned model and tokenizer
MODEL_PATH = "ktr008/sentiment"
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH,timeout=600)
config = AutoConfig.from_pretrained(MODEL_PATH)

# Preprocess function
def preprocess(text):
    new_text = []
    for t in text.split(" "):
        t = '@user' if t.startswith('@') and len(t) > 1 else t
        t = 'http' if t.startswith('http') else t
        new_text.append(t)
    return " ".join(new_text)

# Prediction function
def predict_sentiment(text):
    text = preprocess(text)
    encoded_input = tokenizer(text, return_tensors='pt')
    output = model(**encoded_input)
    scores = output[0][0].detach().numpy()
    scores = softmax(scores)

    # Get sentiment labels and scores (formatted to 4 decimal places)
    ranking = np.argsort(scores)[::-1]
    output_text = "\n".join(
        [f"{config.id2label[ranking[i]]}: {float(scores[ranking[i]]):.4f}" for i in range(scores.shape[0])]
    )
    
    return output_text

# Gradio Interface
interface = gr.Interface(
    fn=predict_sentiment,
    inputs=gr.Textbox(lines=3, placeholder="Enter text..."),
    outputs=gr.Textbox(),  # Display output as plain text
    title="Sentiment Analysis App",
    description="Enter a sentence to analyze its sentiment (Positive, Neutral, Negative).",
)

# Launch the app
interface.launch()