from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import gradio as gr
import numpy as np

# Specify the name of the pretrained sentiment analysis model
MODEL_NAME = "cardiffnlp/twitter-xlm-roberta-base-sentiment"

# Load the tokenizer associated with the model (converts text to tokens)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# Load the pretrained model for sequence classification (sentiment analysis)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

# Labels corresponding to the model's output classes in Georgian language
labels = ['ნეგატიური', 'ნეიტრალური', 'პოზიტიური']

# Define a function to classify sentiment of the input text
def classify_sentiment(text):
    # Tokenize input text and convert to PyTorch tensors, truncate if too long
    inputs = tokenizer(text, return_tensors="pt", truncation=True)

    # Disable gradient calculations for inference
    with torch.no_grad():
        # Pass the tokens through the model to get output logits
        outputs = model(**inputs)
        logits = outputs.logits
        # Convert logits to probabilities using softmax function
        probs = torch.nn.functional.softmax(logits, dim=1).numpy()[0]

    # Find the label with the highest probability
    top_label = labels[np.argmax(probs)]
    confidence = np.max(probs)

    # Return a dictionary with all labels and their probabilities for display
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

# Set up a Gradio interface to interact with the classification function
iface = gr.Interface(
    fn=classify_sentiment,  # function to run when input is given
    inputs=gr.Textbox(lines=3, placeholder="შეიყვანეთ ტვიტი ..."),  # multi-line textbox for input text
    outputs=gr.Label(num_top_classes=3),  # output: show all three sentiment labels with probabilities
    title="Twitter-ის განწყობის კლასიფიკატორი",  # title of the interface in Georgian
    description="იყენებს CardiffNLP-ის მრავალენოვან RoBERTa მოდელს ტვიტების დადებით, ნეიტრალურ ან უარყოფითად კლასიფიცირებისთვის."  # description in Georgian
)

# Launch the Gradio app with public sharing enabled
iface.launch(share=True)