import torch
import gradio as gr
from transformers import DebertaV2Tokenizer
from modeling_multitask_bias import MultiTaskBiasModel

REPO = "piyush333/deberta-v3-multitask-bias-detector-mach-1"
CKPT = "model_dpo_epoch_5.pt"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

label_maps = {
    "political": {0: "left", 1: "neutral", 2: "right"},
    "gender": {0: "misogynist", 1: "neutral", 2: "misandrist"},
    "immigration": {0: "anti", 1: "neutral", 2: "pro"}
}

def load_model():
    tok = DebertaV2Tokenizer.from_pretrained("microsoft/deberta-v3-base")
    mdl = MultiTaskBiasModel()
    state = torch.hub.load_state_dict_from_url(
        f"https://huggingface.co/{REPO}/resolve/main/{CKPT}",
        map_location="cpu"
    )
    mdl.load_state_dict(state)
    mdl.to(DEVICE).eval()
    return tok, mdl

tokenizer, model = load_model()

def predict(text, task):
    enc = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(DEVICE)
    with torch.no_grad():
        logits = model(enc["input_ids"], enc["attention_mask"], [task])
        probs = torch.softmax(logits, dim=1)[0].tolist()
    classes = [label_maps[task][i] for i in range(3)]
    pred = classes[int(torch.tensor(probs).argmax())]
    return pred, {cls: round(float(p), 4) for cls, p in zip(classes, probs)}

demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(label="Text", lines=4, placeholder="Enter text..."),
        gr.Dropdown(["political", "gender", "immigration"], label="Task")
    ],
    outputs=[gr.Textbox(label="Prediction"), gr.JSON(label="Probabilities")],
    title="DeBERTaV3 Multi-Task Bias Detector (Mach-1)",
    description="Detects stance bias in political, gender, and immigration domains"
)

if __name__ == "__main__":
    demo.launch()