| import torch | |
| import numpy as np | |
| from transformers import BertTokenizer | |
| from model import HybridModel | |
| LABELS = ['HS', 'Abusive', 'HS_Individual', 'HS_Group', 'HS_Religion', 'HS_Race', | |
| 'HS_Physical', 'HS_Gender', 'HS_Other', 'HS_Weak', 'HS_Moderate', 'HS_Strong'] | |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| def load_model_and_thresholds(): | |
| model = HybridModel() | |
| model.load_state_dict(torch.load("best_model_dataScrap_final.pt", map_location=DEVICE)) | |
| model.to(DEVICE) | |
| model.eval() | |
| tokenizer = BertTokenizer.from_pretrained("indobenchmark/indobert-base-p1") | |
| thresholds = np.load("optimal_thresholds_dataScrap_final.npy") | |
| return model, tokenizer, thresholds | |
| def predict(text, model, tokenizer, thresholds): | |
| encoding = tokenizer(text, return_tensors='pt', padding='max_length', truncation=True, max_length=128) | |
| input_ids = encoding["input_ids"].to(DEVICE) | |
| attention_mask = encoding["attention_mask"].to(DEVICE) | |
| with torch.no_grad(): | |
| probs = model(input_ids, attention_mask).squeeze(0).cpu().numpy() | |
| return {label: float(prob) for label, prob in zip(LABELS, probs)} |