File size: 6,775 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | """
predict.py — Inference wrapper for the fine-tuned BERT chapter classifier
Usage:
from Classification.predict import ClassifierPredictor
predictor = ClassifierPredictor()
result = predictor.predict("Which organelle is the powerhouse of the cell?")
print(result)
# → {"chapter": "Cell: The Unit of Life", "chapter_id": 8, "confidence": 0.97, "top3": [...]}
"""
import json
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForSequenceClassification
# ==============================
# CONFIG
# ==============================
DEFAULT_MODEL_DIR = "Classification/checkpoints/best_model"
LABEL_MAP_PATH = "Classification/data/label_map.json"
MAX_LENGTH = 128
CONFIDENCE_WARNING = 0.6 # Warn if top prediction confidence is below this
class ClassifierPredictor:
"""
Loads a fine-tuned bert-base-uncased checkpoint and predicts
the NCERT Biology chapter for a given question.
"""
def __init__(self, model_dir: str = DEFAULT_MODEL_DIR):
print(f"[Classifier] Loading model from: {model_dir}")
# Load label map {str(chapter_id): chapter_name}
with open(LABEL_MAP_PATH) as f:
raw_map = json.load(f)
# Keys from JSON are strings — convert to int
self.label_map = {int(k): v for k, v in raw_map.items()}
# Device
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"[Classifier] Using device: {self.device}")
# Load model & tokenizer
self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
self.model = AutoModelForSequenceClassification.from_pretrained(model_dir)
self.model.eval()
self.model.to(self.device)
print(f"[Classifier] Ready — {len(self.label_map)} chapters.")
def predict(self, question: str, top_k: int = 3) -> dict:
"""
Predicts the chapter for a given NEET Biology question.
Returns:
{
"chapter" : "Biomolecules",
"chapter_id" : 9,
"confidence" : 0.94,
"top3" : [
{"chapter": "Biomolecules", "chapter_id": 9, "confidence": 0.94},
{"chapter": "Cell Organelles", "chapter_id": 8, "confidence": 0.04},
{"chapter": "The Living World", "chapter_id": 1, "confidence": 0.01},
]
}
"""
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 # (1, num_labels)
probs = F.softmax(logits, dim=-1).squeeze(0) # (num_labels,)
# Calculate Normalized Entropy
num_labels = len(self.label_map)
entropy = -torch.sum(probs * torch.log(probs + 1e-9)).item()
normalized_entropy = entropy / torch.log(torch.tensor(num_labels, dtype=torch.float32)).item()
# Top-k predictions
top_probs, top_ids = torch.topk(probs, k=min(top_k, num_labels))
top_predictions = []
for prob, label_id in zip(top_probs.tolist(), top_ids.tolist()):
top_predictions.append({
"chapter" : self.label_map.get(label_id, f"Unknown({label_id})"),
"chapter_id" : label_id,
"confidence" : round(prob, 4),
})
best = top_predictions[0]
margin = top_probs[0].item() - top_probs[1].item() if len(top_probs) > 1 else 1.0
# Heuristic Keyword Check (Expanded list)
question_lower = question.lower()
bio_keywords = [
"cell", "dna", "plant", "animal", "human", "body", "blood", "heart",
"gene", "protein", "virus", "bacteria", "disease", "reproduction",
"organ", "tissue", "biology", "enzyme", "acid", "carbon", "oxygen",
"water", "bone", "muscle", "lung", "brain", "leaf", "root", "stem",
"flower", "seed", "fruit", "chromosome", "rna", "atp", "metabolism",
"photosynthesis", "respiration", "digestion", "nervous", "hormone"
]
has_bio_signal = any(kw in question_lower for kw in bio_keywords)
# Dynamic Confidence Threshold
# If the question contains no known biology words, be highly skeptical.
min_confidence = 0.55 if has_bio_signal else 0.85
# OOD Rejection Rules
rejected = False
rejection_reason = ""
if best["confidence"] < min_confidence:
rejected = True
if not has_bio_signal:
rejection_reason = f"No biological keywords found & Confidence too low ({best['confidence']:.2f} < 0.85 for generic text)"
else:
rejection_reason = f"Top-1 confidence too low ({best['confidence']:.2f} < 0.55)"
elif margin < 0.12:
rejected = True
rejection_reason = f"Margin between top 2 chapters too small ({margin:.2f} < 0.12)"
elif normalized_entropy > 0.86:
rejected = True
rejection_reason = f"Model uncertainty too high (Entropy {normalized_entropy:.2f} > 0.86)"
if best["confidence"] < CONFIDENCE_WARNING and not rejected:
print(
f"[Classifier] ⚠️ Low confidence ({best['confidence']:.2f}) for question:\n"
f" '{question}'\n"
f" Consider checking if this chapter is in the training data."
)
return {
"chapter" : best["chapter"],
"chapter_id" : best["chapter_id"],
"confidence" : best["confidence"],
"top3" : top_predictions,
"rejected" : rejected,
"rejection_reason": rejection_reason,
"entropy" : round(normalized_entropy, 4),
"margin" : round(margin, 4)
}
# ==============================
# QUICK TEST
# ==============================
if __name__ == "__main__":
predictor = ClassifierPredictor()
test_questions = [
"What is the role of mitochondria in ATP synthesis?",
"Explain the process of endospore formation in bacteria.",
"What is the significance of the two-kingdom classification?",
]
for q in test_questions:
result = predictor.predict(q)
print(f"\nQ: {q}")
print(f"→ Chapter: {result['chapter']} (confidence: {result['confidence']:.2%})")
top3_display = [(p['chapter'], f"{p['confidence']:.2%}") for p in result['top3']]
print(f" Top 3: {top3_display}")
|