src/predictor.py

"""
UI-ready inference wrapper for the two-head complexity model.
"""

from __future__ import annotations

import json
import pickle
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import numpy as np
import torch

from linguistic_features import extract_linguistic_features
from model_loading import load_model_from_checkpoint
from utils import (
    EXPORT_DIR,
    HARD_LEVELS,
    ID_TO_LEVEL,
    ID_TO_REASON,
    LEVEL_ORDER,
    MODELS,
    REASON_ORDER,
    difficult_class_probability,
    tokenize_lcp_input,
)


@dataclass
class PredictionResult:
    complexity_level: str
    level_id: int
    level_probs: dict[str, float]
    difficult_class_prob: float  # P(Hard)+P(Very Hard) from 5-class softmax (auxiliary)
    reason: str | None
    reason_id: int | None
    reason_probs: dict[str, float] | None

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


class ComplexityPredictor:
    """Load once, call predict() from your UI or API."""

    def __init__(
        self,
        model,
        tokenizer,
        meta: dict,
        max_length: int = 192,
        device: str | None = None,
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.meta = meta
        self.model_key = meta["model_key"]
        self.encoding = meta.get("encoding", "span_mark")
        self.max_length = max_length
        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        self.model.to(self.device)
        self.model.eval()

    @classmethod
    def from_checkpoint(cls, checkpoint_path: str | Path, device: str | None = None) -> "ComplexityPredictor":
        dev = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        model, tokenizer, meta = load_model_from_checkpoint(checkpoint_path, dev)
        return cls(model, tokenizer, meta, device=str(dev))

    @classmethod
    def load(cls, export_dir: str | Path, device: str | None = None) -> "ComplexityPredictor":
        export_dir = Path(export_dir)
        with open(export_dir / "config.json") as f:
            config = json.load(f)

        from transformers import AutoTokenizer
        from two_head_model import TwoHeadModel

        model_key = config["model_key"]
        tokenizer = AutoTokenizer.from_pretrained(export_dir / "tokenizer")
        model = TwoHeadModel(
            model_name=MODELS[model_key],
            pooling_mode=config.get("pooling", "span"),
            use_linguistic_features=config.get("use_linguistic_features", False),
        )
        model.encoder.resize_token_embeddings(len(tokenizer))
        weights = torch.load(export_dir / "model_weights.pt", map_location="cpu", weights_only=True)
        model.load_state_dict(weights)
        model.set_tgt_token_ids(config["tgt_id"], config.get("tgt_end_id"))

        meta = {
            "model_key": model_key,
            "encoding": config.get("encoding", "span_mark"),
            "pooling": config.get("pooling", "span"),
            "use_linguistic_features": config.get("use_linguistic_features", False),
            "tgt_id": config["tgt_id"],
            "tgt_end_id": config.get("tgt_end_id"),
            "feat_mean": np.array(config["feat_mean"]) if config.get("feat_mean") else None,
            "feat_std": np.array(config["feat_std"]) if config.get("feat_std") else None,
        }
        return cls(model, tokenizer, meta, max_length=config.get("max_length", 192), device=device)

    def save(self, export_dir: str | Path) -> Path:
        export_dir = Path(export_dir)
        export_dir.mkdir(parents=True, exist_ok=True)

        config = {
            "model_key": self.model_key,
            "model_name": MODELS[self.model_key],
            "tgt_id": self.meta.get("tgt_id"),
            "tgt_end_id": self.meta.get("tgt_end_id"),
            "encoding": self.encoding,
            "pooling": self.meta.get("pooling", "span"),
            "use_linguistic_features": self.meta.get("use_linguistic_features", False),
            "feat_mean": self.meta.get("feat_mean").tolist() if self.meta.get("feat_mean") is not None else None,
            "feat_std": self.meta.get("feat_std").tolist() if self.meta.get("feat_std") is not None else None,
            "max_length": self.max_length,
            "level_labels": LEVEL_ORDER,
            "reason_labels": REASON_ORDER,
            "hard_levels": list(HARD_LEVELS),
            "inference_rule": "Reason is returned only when predicted level is Hard or Very Hard.",
            "task": "5_class_complexity_prediction",
            "outputs": "discrete_level_and_optional_reason_class",
        }
        with open(export_dir / "config.json", "w") as f:
            json.dump(config, f, indent=2)

        torch.save(self.model.state_dict(), export_dir / "model_weights.pt")
        self.tokenizer.save_pretrained(export_dir / "tokenizer")
        return export_dir

    def save_pickle(self, pickle_path: str | Path) -> Path:
        pickle_path = Path(pickle_path)
        pickle_path.parent.mkdir(parents=True, exist_ok=True)
        with open(pickle_path, "wb") as f:
            pickle.dump(self, f)
        return pickle_path

    @classmethod
    def load_pickle(cls, pickle_path: str | Path, device: str | None = None) -> "ComplexityPredictor":
        with open(pickle_path, "rb") as f:
            obj = pickle.load(f)
        if not isinstance(obj, cls):
            raise TypeError(f"Expected ComplexityPredictor, got {type(obj)}")
        if device:
            obj.device = torch.device(device)
            obj.model.to(obj.device)
        obj.model.eval()
        return obj

    @torch.no_grad()
    def predict(self, sentence: str, target_word: str, corpus: str | None = None) -> PredictionResult:
        enc = tokenize_lcp_input(
            self.tokenizer,
            sentence,
            target_word,
            encoding=self.encoding,
            max_length=self.max_length,
            corpus=corpus,
        )
        enc = {k: v.to(self.device) for k, v in enc.items()}

        kwargs = {"input_ids": enc["input_ids"], "attention_mask": enc["attention_mask"]}
        if self.meta.get("use_linguistic_features"):
            feats = extract_linguistic_features(sentence, target_word)
            if self.meta.get("feat_mean") is not None:
                feats = (feats - self.meta["feat_mean"]) / self.meta["feat_std"]
            kwargs["linguistic_features"] = torch.tensor(feats, dtype=torch.float32).unsqueeze(0).to(self.device)

        out = self.model(**kwargs)
        level_probs_t = torch.softmax(out.level_logits, dim=-1)[0].cpu()
        reason_probs_t = torch.softmax(out.reason_logits, dim=-1)[0].cpu()

        level_id = int(level_probs_t.argmax())
        level_label = ID_TO_LEVEL[level_id]
        level_probs = {ID_TO_LEVEL[i]: float(level_probs_t[i]) for i in range(5)}
        diff_prob = difficult_class_probability(level_probs_t)

        reason_label = None
        reason_id = None
        reason_probs = None
        if level_label in HARD_LEVELS:
            reason_id = int(reason_probs_t.argmax())
            reason_label = ID_TO_REASON[reason_id]
            reason_probs = {ID_TO_REASON[i]: float(reason_probs_t[i]) for i in range(3)}

        return PredictionResult(
            complexity_level=level_label,
            level_id=level_id,
            level_probs=level_probs,
            difficult_class_prob=diff_prob,
            reason=reason_label,
            reason_id=reason_id,
            reason_probs=reason_probs,
        )

    def predict_batch(self, items: list[dict[str, str]]) -> list[PredictionResult]:
        return [self.predict(item["sentence"], item["target_word"], item.get("corpus")) for item in items]