Spaces:

PranilB
/

Question_Answering_Classification

Sleeping

Question_Answering_Classification / Classification /predict.py

Pranil Bankar

Deploy Application

4d67ec7 5 months ago

6.78 kB

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

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