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

MODEL_REPO = "RobroKools/vad-bert"

tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_REPO)
model.eval()

def predict_vad(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
    with torch.no_grad():
        outputs = model(**inputs)
    vad_scores = outputs.logits[0].numpy()
    valence, arousal, dominance = float(vad_scores[0]), float(vad_scores[1]), float(vad_scores[2])
    return valence, arousal, dominance

iface = gr.Interface(
    fn=predict_vad,
    inputs=gr.Textbox(lines=4, placeholder="Enter your journal text here..."),
    outputs=[
        gr.Number(label="Valence"),
        gr.Number(label="Arousal"),
        gr.Number(label="Dominance"),
    ],
    title="VAD Emotion Predictor",
    description="Enter text to get Valence, Arousal, and Dominance scores from a fine-tuned BERT model.",
)

iface.launch()