File size: 2,478 Bytes
4d67ec7 | 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 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | """
predict_multilabel.py — Inference wrapper for fine-tuned Multi-Label BERT
"""
import json
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
DEFAULT_MODEL_DIR = "Classification/checkpoints_multilabel/best_model"
LABEL_MAP_PATH = "Classification/data/label_map.json"
MAX_LENGTH = 128
THRESHOLD = 0.5
class MultiLabelPredictor:
def __init__(self, model_dir: str = DEFAULT_MODEL_DIR):
print(f"[MultiLabel Predictor] Loading model from: {model_dir}")
with open(LABEL_MAP_PATH) as f:
raw_map = json.load(f)
self.label_map = {int(k): v for k, v in raw_map.items()}
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"[MultiLabel Predictor] Using device: {self.device}")
self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
self.model = AutoModelForSequenceClassification.from_pretrained(model_dir)
self.model.eval()
self.model.to(self.device)
def predict(self, question: str) -> list:
inputs = self.tokenizer(
question,
max_length=MAX_LENGTH,
padding="max_length",
truncation=True,
return_tensors="pt",
)
inputs = {k: v.to(self.device) for k, v in inputs.items()}
with torch.no_grad():
logits = self.model(**inputs).logits
# Apply sigmoid for independent probabilities
probs = torch.sigmoid(logits).squeeze(0)
results = []
for label_id, prob in enumerate(probs.tolist()):
if prob >= THRESHOLD:
results.append({
"chapter": self.label_map.get(label_id, f"Unknown({label_id})"),
"chapter_id": label_id,
"confidence": round(prob, 4),
})
# Sort by confidence
results = sorted(results, key=lambda x: x["confidence"], reverse=True)
return results
if __name__ == "__main__":
import sys
try:
predictor = MultiLabelPredictor()
except Exception as e:
print(f"Error loading model: {e}")
sys.exit(1)
qs = ["Which structures are common to both plant cells and animal cells?"]
for q in qs:
print(f"\nQ: {q}")
preds = predictor.predict(q)
print("Predictions:")
for p in preds:
print(f" - {p['chapter']} : {p['confidence']:.2%}")
|