admesh/agentic-intent-classifier

1from __future__ import annotations2 3from functools import lru_cache4 5import torch6 7try:8    from .config import IAB_PARENT_FALLBACK_CONFIDENCE_FLOOR, _looks_like_local_hf_model_dir  # type: ignore9    from .iab_taxonomy import get_iab_taxonomy, parse_path_label, path_to_label  # type: ignore10    from .model_runtime import get_head  # type: ignore11except ImportError:12    from config import IAB_PARENT_FALLBACK_CONFIDENCE_FLOOR, _looks_like_local_hf_model_dir13    from iab_taxonomy import get_iab_taxonomy, parse_path_label, path_to_label14    from model_runtime import get_head15 16 17def round_score(value: float) -> float:18    return round(float(value), 4)19 20 21@lru_cache(maxsize=1)22def _prefix_label_ids() -> dict[tuple[str, ...], list[int]]:23    head = get_head("iab_content")24    prefix_map: dict[tuple[str, ...], list[int]] = {}25    for label, label_id in head.config.label2id.items():26        path = parse_path_label(label)27        for depth in range(1, len(path) + 1):28            prefix_map.setdefault(path[:depth], []).append(label_id)29    return prefix_map30 31 32def _effective_exact_threshold(confidence_threshold: float | None) -> float:33    head = get_head("iab_content")34    if confidence_threshold is None:35        return float(head.calibration.confidence_threshold)36    return min(max(float(confidence_threshold), 0.0), 1.0)37 38 39def _effective_parent_threshold(exact_threshold: float) -> float:40    return min(max(IAB_PARENT_FALLBACK_CONFIDENCE_FLOOR, exact_threshold), 1.0)41 42 43def _build_prediction(44    accepted_path: tuple[str, ...],45    *,46    exact_label: str,47    confidence: float,48    raw_confidence: float,49    exact_threshold: float,50    calibrated: bool,51    meets_confidence_threshold: bool,52    mapping_mode: str,53    stopped_reason: str,54) -> dict:55    taxonomy = get_iab_taxonomy()56    return {57        "label": path_to_label(accepted_path),58        "exact_label": exact_label,59        "path": list(accepted_path),60        "confidence": round_score(confidence),61        "raw_confidence": round_score(raw_confidence),62        "confidence_threshold": round_score(exact_threshold),63        "calibrated": calibrated,64        "meets_confidence_threshold": meets_confidence_threshold,65        "content": taxonomy.build_content_object(66            accepted_path,67            mapping_mode=mapping_mode,68            mapping_confidence=confidence,69        ),70        "mapping_mode": mapping_mode,71        "mapping_confidence": round_score(confidence),72        "source": "supervised_classifier",73        "stopped_reason": stopped_reason,74    }75 76 77def predict_iab_content_classifier_batch(78    texts: list[str],79    confidence_threshold: float | None = None,80) -> list[dict | None]:81    if not texts:82        return []83 84    head = get_head("iab_content")85    # `SequenceClassifierHead` will raise if the folder exists but is incomplete86    # (missing `model.safetensors` / `pytorch_model.bin`). Treat that as "no model".87    if not _looks_like_local_hf_model_dir(head.config.model_dir):88        return [None for _ in texts]89 90    raw_probs, calibrated_probs = head.predict_probs_batch(texts)91    prefix_map = _prefix_label_ids()92    exact_threshold = _effective_exact_threshold(confidence_threshold)93    parent_threshold = _effective_parent_threshold(exact_threshold)94    predictions: list[dict | None] = []95 96    for raw_row, calibrated_row in zip(raw_probs, calibrated_probs):97        pred_id = int(torch.argmax(calibrated_row).item())98        exact_label = head.model.config.id2label[pred_id]99        exact_path = parse_path_label(exact_label)100        exact_confidence = float(calibrated_row[pred_id].item())101        exact_raw_confidence = float(raw_row[pred_id].item())102 103        if exact_confidence >= exact_threshold:104            predictions.append(105                _build_prediction(106                    exact_path,107                    exact_label=exact_label,108                    confidence=exact_confidence,109                    raw_confidence=exact_raw_confidence,110                    exact_threshold=exact_threshold,111                    calibrated=head.calibration.calibrated,112                    meets_confidence_threshold=True,113                    mapping_mode="exact",114                    stopped_reason="exact_threshold_met",115                )116            )117            continue118 119        accepted_path = exact_path[:1]120        accepted_confidence = float(calibrated_row[prefix_map[accepted_path]].sum().item())121        accepted_raw_confidence = float(raw_row[prefix_map[accepted_path]].sum().item())122        meets_confidence_threshold = False123        stopped_reason = "top_level_safe_fallback"124 125        for depth in range(len(exact_path) - 1, 0, -1):126            prefix = exact_path[:depth]127            prefix_ids = prefix_map[prefix]128            prefix_confidence = float(calibrated_row[prefix_ids].sum().item())129            prefix_raw_confidence = float(raw_row[prefix_ids].sum().item())130            if prefix_confidence >= parent_threshold:131                accepted_path = prefix132                accepted_confidence = prefix_confidence133                accepted_raw_confidence = prefix_raw_confidence134                meets_confidence_threshold = True135                stopped_reason = "parent_fallback_threshold_met"136                break137 138        predictions.append(139            _build_prediction(140                accepted_path,141                exact_label=exact_label,142                confidence=accepted_confidence,143                raw_confidence=accepted_raw_confidence,144                exact_threshold=exact_threshold,145                calibrated=head.calibration.calibrated,146                meets_confidence_threshold=meets_confidence_threshold,147                mapping_mode="nearest_equivalent",148                stopped_reason=stopped_reason,149            )150        )151 152    return predictions153 154 155def predict_iab_content_classifier(text: str, confidence_threshold: float | None = None) -> dict | None:156    return predict_iab_content_classifier_batch([text], confidence_threshold=confidence_threshold)[0]157