BaseModelGBert / how_to_load.py
AnnNaserNabil's picture
Upload 10 files
d332465 verified
# How to load this model:
from transformers import AutoModel, AutoTokenizer
import torch
import torch.nn as nn
import json
# Load encoder
encoder = AutoModel.from_pretrained("./outputs/final_baseline_best")
# Load classifier config
with open("./outputs/final_baseline_best/classifier_config.json", 'r') as f:
c_config = json.load(f)
num_labels = c_config.get('num_labels', 1)
hidden_size = c_config.get('hidden_size', 768)
# Reconstruct classifier
classifier = nn.Sequential(
nn.Linear(hidden_size, 256),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(256, num_labels)
)
classifier.load_state_dict(torch.load("./outputs/final_baseline_best/classifier.pt"))
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("./outputs/final_baseline_best")
# Inference function
def predict(text):
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
with torch.no_grad():
outputs = encoder(**inputs)
cls_embedding = outputs.last_hidden_state[:, 0, :]
logits = classifier(cls_embedding)
probs = torch.sigmoid(logits)
return probs.item()
# Example
text = "আপনার বাংলা টেক্সট এখানে"
prob = predict(text)
print(f"Hate Speech Probability: {prob:.4f}")
print(f"Prediction: {'Hate Speech' if prob > 0.5 else 'Non-Hate Speech'}")