import torch
from transformers import BertTokenizer
import numpy as np
from model import HybridModel

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
LABELS = ['HS_Individual', 'HS_Group', 'HS_Religion', 'HS_Race', 'HS_Physical',
          'HS_Gender', 'HS_Other', 'HS_Weak', 'HS_Moderate', 'HS_Strong']
THRESHOLDS = [0.45, 0.44, 0.43, 0.42, 0.46, 0.44, 0.47, 0.41, 0.40, 0.43]

def load_model():
    model = HybridModel()
    model.load_state_dict(torch.load("best_model.pt", map_location=DEVICE))
    model.to(DEVICE)
    model.eval()
    tokenizer = BertTokenizer.from_pretrained("indobenchmark/indobert-base-p1")
    return model, tokenizer

def predict(text, model, tokenizer):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
    input_ids = inputs["input_ids"].to(DEVICE)
    attention_mask = inputs["attention_mask"].to(DEVICE)

    with torch.no_grad():
        outputs = model(input_ids, attention_mask)
    probs = outputs.cpu().numpy()[0]

    # Buat dictionary {label: prob} untuk Gradio Label with confidence
    result = {label: float(prob) for label, prob in zip(LABELS, probs)}
    return result