File size: 1,182 Bytes
0b32a79 dc14b12 0b32a79 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
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 |