import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from huggingface_hub import hf_hub_download

REPO_ID  = "SentilyticsOPJ/bertic_model" 
FILENAME = "bertic_model.pt"

path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
obj  = torch.load(path, map_location="cpu", weights_only=False)

hf_model_name = obj["hf_model_name"]     
labels        = obj["label_classes"]     
max_len       = obj["max_len"]           

tokenizer = AutoTokenizer.from_pretrained(hf_model_name)
model = AutoModelForSequenceClassification.from_pretrained(
    hf_model_name,
    num_labels=len(labels),
)

model.load_state_dict(obj["state_dict"])
model.eval()


def predict(text):
    if not text.strip():
        return ""
    enc = tokenizer(
        text,
        max_length=max_len,
        padding="max_length",
        truncation=True,
        return_tensors="pt",
    )
    with torch.no_grad():
        logits = model(
            input_ids=enc["input_ids"],
            attention_mask=enc["attention_mask"],
        ).logits
    return labels[int(logits.argmax(1))]


demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=3, placeholder="Unesi tekst...", label="Text"),
    outputs=gr.Textbox(label="Sentiment"),
    title="Sentiment Analysis (BERTić, fine-tuned)",
    description="Five-class sentiment: mixed, negative, neutral, positive, sarcastic.",
    examples=["Volim kavu", "Ovaj doktor je loš", "Dan je bio ok"],
)

demo.launch()