import os

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

MODEL_ID = os.getenv("MODEL_ID", "Pectics/vad-macbert")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID).to(DEVICE)
model.eval()


def predict(text):
    if not text or not text.strip():
        return {"error": "Input text is empty."}

    inputs = tokenizer(
        text,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=128,
    )
    inputs = {key: value.to(DEVICE) for key, value in inputs.items()}
    with torch.no_grad():
        outputs = model(**inputs)
    vad = outputs.logits.squeeze().tolist()
    return {"vad": vad}


with gr.Blocks(title="Vad MacBERT Demo") as demo:
    gr.Markdown(
        "\n".join(
            [
                "# Pectics/vad-macbert Demo",
                f"Model: `{MODEL_ID}`",
                "中文输入，返回序列分类头的 VAD 分数。",
            ]
        )
    )
    text_input = gr.Textbox(
        label="Input text",
        lines=4,
        placeholder="请输入中文文本。",
    )
    run_btn = gr.Button("Run")
    output = gr.JSON(label="VAD")

    run_btn.click(fn=predict, inputs=text_input, outputs=output)

    gr.Examples(
        examples=[
            "这部电影让我很感动。",
            "我对这个结果很失望。",
        ],
        inputs=text_input,
    )

demo.launch()