|
|
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] |
|
|
|
|
|
|
|
|
result = {label: float(prob) for label, prob in zip(LABELS, probs)} |
|
|
return result |
