| """ |
| 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 |
|
|
| |
| 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), |
| }) |
| |
| |
| 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%}") |
|
|
