import gradio as gr
import torch
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model import TRALSem
from transformers import BertTokenizer

CHECKPOINT = "tralsem_imdb_freelb_best.pt"
NUM_LABELS = 2
MAX_LEN = 128

tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = TRALSem(num_labels=NUM_LABELS)
ckpt = torch.load(CHECKPOINT, map_location="cpu", weights_only=False)
model.load_state_dict(ckpt["model_state_dict"])
model.eval()

LABELS = {0: "Negative", 1: "Positive"}

def predict(text):
    enc = tokenizer(
        text,
        max_length=MAX_LEN,
        padding="max_length",
        truncation=True,
        return_tensors="pt",
        return_token_type_ids=True
    )
    with torch.no_grad():
        logits = model(
            enc["input_ids"],
            enc["attention_mask"],
            enc["token_type_ids"]
        )
    import torch.nn.functional as F
    probs = F.softmax(logits, dim=-1).squeeze(0).tolist()
    return {LABELS[i]: round(probs[i], 4) for i in range(NUM_LABELS)}

demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=5, placeholder="Enter text to analyse sentiment..."),
    outputs=gr.Label(num_top_classes=2),
    title="A Robust Transformer Based Sentiment Analysis System",
    description="TRALSem — Adversarially trained (FreeLB) on IMDB dataset",
    examples=[
        ["This movie was absolutely brilliant!"],
        ["Terrible experience, waste of time."],
        ["It was okay, nothing special."]
    ]
)
demo.launch()