admesh/agentic-intent-classifier

1from __future__ import annotations2 3import json4import re5from functools import lru_cache6 7import torch8import torch.nn.functional as F9from transformers import AutoModel, AutoTokenizer10 11from config import (12    IAB_RETRIEVAL_DEPTH_BONUS,13    IAB_RETRIEVAL_MODEL_MAX_LENGTH,14    IAB_RETRIEVAL_MODEL_NAME,15    IAB_RETRIEVAL_PREFIX_CONFIDENCE_THRESHOLDS,16    IAB_RETRIEVAL_TOP_K,17    IAB_TAXONOMY_EMBEDDINGS_PATH,18    IAB_TAXONOMY_NODES_PATH,19    IAB_TAXONOMY_VERSION,20    ensure_artifact_dirs,21)22from iab_taxonomy import IabNode, get_iab_taxonomy, path_to_label23 24RETRIEVAL_STOPWORDS = {25    "a",26    "an",27    "and",28    "are",29    "as",30    "at",31    "best",32    "buy",33    "for",34    "from",35    "how",36    "i",37    "in",38    "is",39    "it",40    "me",41    "my",42    "need",43    "of",44    "on",45    "or",46    "should",47    "the",48    "to",49    "tonight",50    "what",51    "which",52    "with",53}54GTE_QWEN_QUERY_INSTRUCTION = "Given a user query, retrieve the most relevant IAB content taxonomy category."55 56 57def round_score(value: float) -> float:58    return round(float(value), 4)59 60 61def _normalize_keyword(value: str) -> str:62    value = value.lower().replace("&", " and ")63    value = re.sub(r"[^a-z0-9]+", " ", value)64    return " ".join(value.split())65 66 67def _keyword_tokens(value: str) -> set[str]:68    return {69        token70        for token in _normalize_keyword(value).split()71        if token and token not in RETRIEVAL_STOPWORDS and len(token) > 172    }73 74 75def _is_gte_qwen_model(model_name: str) -> bool:76    normalized = model_name.lower()77    return "gte-qwen" in normalized78 79 80def _last_token_pool(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:81    left_padding = bool(torch.all(attention_mask[:, -1] == 1))82    if left_padding:83        return last_hidden_state[:, -1]84 85    sequence_lengths = attention_mask.sum(dim=1) - 186    batch_indices = torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device)87    return last_hidden_state[batch_indices, sequence_lengths]88 89 90def _node_keywords(node: IabNode) -> list[str]:91    keywords = {node.label, node.path_label}92    keywords.update(node.path)93    normalized = {_normalize_keyword(keyword) for keyword in keywords if keyword.strip()}94    return sorted(keyword for keyword in normalized if keyword)95 96 97def _node_retrieval_text(node: IabNode) -> str:98    keywords = _node_keywords(node)99    parts = [100        f"IAB category path: {node.path_label}",101        f"Canonical label: {node.label}",102        f"Tier depth: {node.level}",103    ]104    if len(node.path) > 1:105        parts.append(f"Parent path: {' > '.join(node.path[:-1])}")106    if keywords:107        parts.append(f"Keywords: {', '.join(keywords)}")108    return ". ".join(parts)109 110 111def _serialize_node(node: IabNode) -> dict:112    return {113        "unique_id": node.unique_id,114        "parent_id": node.parent_id,115        "label": node.label,116        "path": list(node.path),117        "path_label": node.path_label,118        "level": node.level,119        "keywords": _node_keywords(node),120        "retrieval_text": _node_retrieval_text(node),121    }122 123 124class LocalTextEmbedder:125    def __init__(self, model_name: str, max_length: int):126        self.model_name = model_name127        self.max_length = max_length128        self._tokenizer = None129        self._model = None130        self._batch_size = 32131        self._device = "cuda" if torch.cuda.is_available() else "cpu"132        self._is_gte_qwen = _is_gte_qwen_model(model_name)133 134    @property135    def tokenizer(self):136        if self._tokenizer is None:137            self._tokenizer = AutoTokenizer.from_pretrained(138                self.model_name,139                trust_remote_code=self._is_gte_qwen,140            )141        return self._tokenizer142 143    @property144    def model(self):145        if self._model is None:146            model_kwargs = {"trust_remote_code": self._is_gte_qwen}147            if self._device == "cuda":148                model_kwargs["torch_dtype"] = torch.float16149            self._model = AutoModel.from_pretrained(self.model_name, **model_kwargs)150            self._model.to(self._device)151            self._model.eval()152        return self._model153 154    def encode_documents(self, texts: list[str], batch_size: int | None = None) -> torch.Tensor:155        return self._encode_texts(texts, batch_size=batch_size, treat_as_query=False)156 157    def encode_queries(self, texts: list[str], batch_size: int | None = None) -> torch.Tensor:158        return self._encode_texts(texts, batch_size=batch_size, treat_as_query=True)159 160    def _encode_texts(161        self,162        texts: list[str],163        batch_size: int | None = None,164        treat_as_query: bool = False,165    ) -> torch.Tensor:166        if not texts:167            return torch.empty(0, 0)168 169        effective_batch_size = batch_size or self._batch_size170        rows: list[torch.Tensor] = []171        for start in range(0, len(texts), effective_batch_size):172            batch_texts = texts[start : start + effective_batch_size]173            if treat_as_query and self._is_gte_qwen:174                batch_texts = [175                    f"Instruct: {GTE_QWEN_QUERY_INSTRUCTION}\nQuery: {text}"176                    for text in batch_texts177                ]178            inputs = self.tokenizer(179                batch_texts,180                return_tensors="pt",181                truncation=True,182                padding=True,183                max_length=self.max_length,184            )185            inputs = {key: value.to(self._device) for key, value in inputs.items()}186            with torch.no_grad():187                outputs = self.model(**inputs)188            hidden = outputs.last_hidden_state189            if self._is_gte_qwen:190                pooled = _last_token_pool(hidden, inputs["attention_mask"])191            else:192                mask = inputs["attention_mask"].unsqueeze(-1)193                pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)194            rows.append(F.normalize(pooled.float(), p=2, dim=1).cpu())195        return torch.cat(rows, dim=0)196 197 198@lru_cache(maxsize=1)199def get_iab_text_embedder() -> LocalTextEmbedder:200    return LocalTextEmbedder(IAB_RETRIEVAL_MODEL_NAME, IAB_RETRIEVAL_MODEL_MAX_LENGTH)201 202 203def build_iab_taxonomy_embedding_index(batch_size: int = 32) -> dict:204    ensure_artifact_dirs()205    taxonomy = get_iab_taxonomy()206    nodes = [_serialize_node(node) for node in taxonomy.nodes]207    embedder = get_iab_text_embedder()208    embeddings = embedder.encode_documents([node["retrieval_text"] for node in nodes], batch_size=batch_size)209 210    IAB_TAXONOMY_NODES_PATH.write_text(json.dumps(nodes, indent=2, sort_keys=True) + "\n", encoding="utf-8")211    torch.save(212        {213            "model_name": embedder.model_name,214            "taxonomy_version": IAB_TAXONOMY_VERSION,215            "embedding_dim": int(embeddings.shape[1]),216            "node_count": len(nodes),217            "embeddings": embeddings,218        },219        IAB_TAXONOMY_EMBEDDINGS_PATH,220    )221    return {222        "taxonomy_version": IAB_TAXONOMY_VERSION,223        "model_name": embedder.model_name,224        "node_count": len(nodes),225        "embedding_dim": int(embeddings.shape[1]),226        "nodes_path": str(IAB_TAXONOMY_NODES_PATH),227        "embeddings_path": str(IAB_TAXONOMY_EMBEDDINGS_PATH),228    }229 230 231class IabEmbeddingRetriever:232    def __init__(self):233        self.taxonomy = get_iab_taxonomy()234        self.embedder = get_iab_text_embedder()235        self._nodes: list[dict] | None = None236        self._embeddings: torch.Tensor | None = None237 238    def _load_index(self) -> bool:239        if self._nodes is not None and self._embeddings is not None:240            return True241        if not IAB_TAXONOMY_NODES_PATH.exists() or not IAB_TAXONOMY_EMBEDDINGS_PATH.exists():242            return False243 244        nodes = json.loads(IAB_TAXONOMY_NODES_PATH.read_text(encoding="utf-8"))245        payload = torch.load(IAB_TAXONOMY_EMBEDDINGS_PATH, map_location="cpu")246        if payload.get("model_name") != IAB_RETRIEVAL_MODEL_NAME:247            return False248        if payload.get("taxonomy_version") != IAB_TAXONOMY_VERSION:249            return False250 251        embeddings = payload.get("embeddings")252        if not isinstance(embeddings, torch.Tensor):253            embeddings = torch.tensor(embeddings, dtype=torch.float32)254        if len(nodes) != embeddings.shape[0]:255            return False256 257        self._nodes = nodes258        self._embeddings = F.normalize(embeddings.float(), p=2, dim=1)259        return True260 261    def ready(self) -> bool:262        return self._load_index()263 264    @staticmethod265    def _score_to_confidence(score: float) -> float:266        return min(max((score + 1.0) / 2.0, 0.0), 1.0)267 268    def _candidate_from_index(self, score: float, index: int) -> dict:269        assert self._nodes is not None270        node = self._nodes[index]271        confidence = self._score_to_confidence(float(score))272        adjusted_confidence = confidence + (IAB_RETRIEVAL_DEPTH_BONUS * max(int(node["level"]) - 1, 0))273        return {274            "unique_id": node["unique_id"],275            "label": node["label"],276            "path": tuple(node["path"]),277            "path_label": node["path_label"],278            "level": int(node["level"]),279            "confidence": round_score(confidence),280            "adjusted_confidence": round_score(adjusted_confidence),281            "keywords": list(node.get("keywords", [])),282        }283 284    def _rerank_candidates(self, query_text: str, candidates: list[dict]) -> list[dict]:285        if not candidates:286            return []287 288        query_normalized = _normalize_keyword(query_text)289        query_tokens = _keyword_tokens(query_text)290 291        reranked = []292        for candidate in candidates:293            keyword_tokens = set()294            for keyword in candidate.get("keywords", []):295                keyword_tokens.update(_keyword_tokens(keyword))296 297            token_overlap = len(query_tokens & keyword_tokens)298            path_overlap = len(query_tokens & _keyword_tokens(candidate["path_label"]))299            lexical_bonus = min(0.04, (0.008 * token_overlap) + (0.004 * path_overlap))300 301            reranked.append(302                {303                    **candidate,304                    "token_overlap": token_overlap,305                    "path_overlap": path_overlap,306                    "lexical_bonus": round_score(lexical_bonus),307                    "rerank_score": round_score(candidate["adjusted_confidence"] + lexical_bonus),308                }309            )310 311        reranked.sort(312            key=lambda item: (313                item["rerank_score"],314                item["adjusted_confidence"],315                item["confidence"],316            ),317            reverse=True,318        )319        return reranked320 321    def _top_candidates_from_embedding(self, query_text: str, query_embedding: torch.Tensor) -> list[dict]:322        if not self._load_index():323            return []324 325        assert self._embeddings is not None326 327        scores = torch.mv(self._embeddings, query_embedding)328        top_k = min(IAB_RETRIEVAL_TOP_K, scores.shape[0])329        top_scores, top_indices = torch.topk(scores, k=top_k)330 331        candidates = [self._candidate_from_index(score, index) for score, index in zip(top_scores.tolist(), top_indices.tolist())]332        return self._rerank_candidates(query_text, candidates)333 334    def _top_candidates(self, text: str) -> list[dict]:335        if not self._load_index():336            return []337        query_embedding = self.embedder.encode_queries([text])[0]338        return self._top_candidates_from_embedding(text, query_embedding)339 340    def _select_path(self, candidates: list[dict]) -> dict | None:341        if not candidates:342            return None343 344        top_candidate = candidates[0]345        top_path = tuple(top_candidate["path"])346        top_margin = round_score(347            top_candidate["confidence"] - candidates[1]["confidence"] if len(candidates) > 1 else top_candidate["confidence"]348        )349        prefix_support: dict[tuple[str, ...], float] = {}350        for depth in range(1, len(top_path) + 1):351            prefix = top_path[:depth]352            prefix_support[prefix] = max(353                candidate["confidence"]354                for candidate in candidates355                if tuple(candidate["path"][:depth]) == prefix356            )357 358        selected_path: tuple[str, ...] | None = None359        selected_threshold = 0.0360        for depth in range(1, len(top_path) + 1):361            threshold = IAB_RETRIEVAL_PREFIX_CONFIDENCE_THRESHOLDS.get(depth, 0.62)362            prefix = top_path[:depth]363            if prefix_support[prefix] >= threshold:364                selected_path = prefix365                selected_threshold = threshold366                continue367            break368 369        if selected_path is None:370            return None371 372        stopped_reason = "accepted" if selected_path == top_path else "parent_fallback"373        if len(top_path) > 1:374            ambiguous_sibling = any(375                tuple(candidate["path"][:-1]) == top_path[:-1]376                and (top_candidate["confidence"] - candidate["confidence"]) <= 0.03377                for candidate in candidates[1:]378            )379            if ambiguous_sibling:380                selected_path = top_path[:-1]381                selected_threshold = IAB_RETRIEVAL_PREFIX_CONFIDENCE_THRESHOLDS.get(len(selected_path), 0.62)382                stopped_reason = "ambiguous_sibling_parent_fallback"383 384        mapping_confidence = prefix_support[selected_path]385        return {386            "path": selected_path,387            "path_label": path_to_label(selected_path),388            "mapping_mode": "nearest_equivalent",389            "mapping_confidence": round_score(mapping_confidence),390            "confidence_threshold": round_score(selected_threshold),391            "top_candidate_confidence": round_score(top_candidate["confidence"]),392            "top_margin": top_margin,393            "stopped_reason": stopped_reason,394        }395 396    def predict(self, text: str) -> dict | None:397        candidates = self._top_candidates(text)398        return self._prediction_from_candidates(candidates)399 400    def _prediction_from_candidates(self, candidates: list[dict]) -> dict | None:401        selection = self._select_path(candidates)402        if selection is None:403            return None404 405        content = self.taxonomy.build_content_object(406            path=selection["path"],407            mapping_mode=selection["mapping_mode"],408            mapping_confidence=selection["mapping_confidence"],409        )410        return {411            "label": selection["path_label"],412            "confidence": selection["mapping_confidence"],413            "raw_confidence": selection["top_candidate_confidence"],414            "confidence_threshold": selection["confidence_threshold"],415            "calibrated": False,416            "meets_confidence_threshold": True,417            "content": content,418            "path": selection["path"],419            "mapping_mode": selection["mapping_mode"],420            "mapping_confidence": selection["mapping_confidence"],421            "source": "embedding_retrieval",422            "retrieval_model_name": IAB_RETRIEVAL_MODEL_NAME,423            "stopped_reason": selection["stopped_reason"],424            "top_margin": selection["top_margin"],425            "top_candidates": [426                {427                    **candidate,428                    "path": list(candidate["path"]),429                    "keywords": candidate["keywords"][:12],430                }431                for candidate in candidates432            ],433        }434 435    def predict_batch(self, texts: list[str], batch_size: int | None = None) -> list[dict | None]:436        if not texts:437            return []438        if not self._load_index():439            return [None for _ in texts]440 441        query_embeddings = self.embedder.encode_queries(texts, batch_size=batch_size)442        return [443            self._prediction_from_candidates(self._top_candidates_from_embedding(text, query_embedding))444            for text, query_embedding in zip(texts, query_embeddings)445        ]446 447 448@lru_cache(maxsize=1)449def get_iab_embedding_retriever() -> IabEmbeddingRetriever:450    return IabEmbeddingRetriever()451 452 453def predict_iab_content_retrieval(text: str) -> dict | None:454    retriever = get_iab_embedding_retriever()455    if not retriever.ready():456        return None457    return retriever.predict(text)458 459 460def predict_iab_content_retrieval_batch(texts: list[str], batch_size: int | None = None) -> list[dict | None]:461    retriever = get_iab_embedding_retriever()462    if not retriever.ready():463        return [None for _ in texts]464    return retriever.predict_batch(texts, batch_size=batch_size)465