Update handler.py

handler.py

@@ -1,6 +1,9 @@
 import os
 import json
-from typing import Any, Dict, List, Optional
+import re
+import unicodedata
+from typing import Any, Dict, List, Optional, Tuple
+from collections import defaultdict
 
 import numpy as np
 import pandas as pd
@@ -8,23 +11,149 @@ import pandas as pd
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-# Optional dependencies (kept in requirements.txt)
 import faiss
 from sentence_transformers import SentenceTransformer, CrossEncoder
 
 
-def softmax(x: np.ndarray) -> np.ndarray:
+# -----------------------------
+# Basic math utils
+# -----------------------------
+def softmax_1d(x: np.ndarray) -> np.ndarray:
     x = x - np.max(x)
     e = np.exp(x)
     return e / np.sum(e)
 
 
+def norm_cos_to_01(cos_sim: float) -> float:
+    # For IndexFlatIP + normalized vectors => cosine in [-1,1], map to [0,1]
+    return float((cos_sim + 1.0) / 2.0)
+
+
+def minmax_norm(vals: List[float]) -> List[float]:
+    if not vals:
+        return vals
+    vmin, vmax = min(vals), max(vals)
+    if abs(vmax - vmin) < 1e-12:
+        return [0.0 for _ in vals]
+    return [(v - vmin) / (vmax - vmin) for v in vals]
+
+
+# ============================================================
+# Cleaning the request
+# ============================================================
+def _strip_after_first_bar(s: str) -> str:
+    """
+    Keep only what is after the first '|', if any.
+    Example:
+      "uom: UN | Fourniture ...\n..." -> "Fourniture ... ..."
+    """
+    if not s or not isinstance(s, str):
+        return ""
+    s = s.replace("\r", " ").replace("\n", " ")
+    if "|" in s:
+        s = s.split("|", 1)[1]
+    s = re.sub(r"\s+", " ", s).strip()
+    return s
+
+
+def fix_encoding(text: str) -> str:
+    """
+    Best-effort fixer for common mojibake: "associés" -> "associés"
+    Works when UTF-8 bytes were decoded as latin1/cp1252.
+    """
+    if not text or not isinstance(text, str):
+        return ""
+
+    # Fast exit: if looks clean, keep
+    # (but still normalize later in clean_custom_code_name)
+    t = text
+
+    # Try latin1 -> utf8
+    try:
+        t2 = t.encode("latin1", errors="strict").decode("utf-8", errors="strict")
+        # Accept only if it actually improves typical patterns
+        if "Ã" in t or "�" in t:
+            return t2
+    except Exception:
+        pass
+
+    # Try cp1252 -> utf8
+    try:
+        t2 = t.encode("cp1252", errors="strict").decode("utf-8", errors="strict")
+        if "Ã" in t or "�" in t:
+            return t2
+    except Exception:
+        pass
+
+    return t
+
+
+def clean_custom_code_name(name: str) -> str:
+    """
+    Cleans a custom code / label text to improve semantic matching.
+
+    Examples:
+      "NPO - 1615 - IT group services re-billing (0006111100)"
+      -> "IT group services re-billing"
+
+      "Fees frais de deplacement associés (0006228000)"
+      -> "Fees frais de deplacement associés"
+
+      "Reception (0006257000)"
+      -> "Reception"
+
+      "1403 - Maintenance matériel informatique (ordinateurs, PC)"
+      -> "Maintenance matériel informatique ordinateurs PC"
+    """
+    if not name or not isinstance(name, str):
+        return ""
+
+    # 1) Fix broken encoding
+    name = fix_encoding(name)
+
+    # 2) Remove numeric codes in parentheses: (0006111100)
+    name = re.sub(r"\(\d{5,}\)", "", name)
+
+    # 3) Remove prefixes like "NPO - 1615 -"
+    name = re.sub(r"^NPO\s*-\s*\d+\s*-\s*", "", name, flags=re.IGNORECASE)
+
+    # 4) Remove numeric prefixes like "1403 -" or "1403 :"
+    name = re.sub(r"^\d+\s*[-:]\s*", "", name)
+
+    # 5) Remove parentheses but keep content
+    name = re.sub(r"[()]", " ", name)
+
+    # 6) Remove square brackets but keep content
+    name = re.sub(r"[\[\]]", " ", name)
+
+    # 7) Remove noisy punctuation (keep letters/digits/spaces)
+    name = re.sub(r"[,;/\\|#@&*=+]", " ", name)
+
+    # Optional: normalize hyphens/long dashes as spaces (helps match)
+    name = re.sub(r"[-–—]", " ", name)
+
+    # 8) Normalize spaces
+    name = re.sub(r"\s+", " ", name).strip()
+
+    # 9) Unicode normalize
+    name = unicodedata.normalize("NFC", name)
+
+    return name
+
+
+# ============================================================
+# Handler
+# ============================================================
 class EndpointHandler:
     """
-    Hugging Face Inference Endpoint Custom Handler
-    - Loads supervised classifier from ./supervised_final
-    - Loads FAISS index + label table for semantic fallback
-    - Optionally loads a reranker (cross-encoder)
+    HF Inference Endpoint Custom Handler (Dual Semantic + Cleaners)
+
+    Artifacts expected in repo root:
+      - supervised_final/  (HF transformers classifier)
+      - faiss_labels.index + label_table.csv            (semantic labels)
+      - faiss_catalog_articles.index + catalog_articles_table.csv (semantic catalog)
+      - biencoder_model_name.txt
+      - optional: rerank_model_name.txt
     """
 
     def __init__(self, path: str = ""):
@@ -33,63 +162,81 @@ class EndpointHandler:
         # ---- 1) Supervised classifier ----
         self.sup_dir = os.path.join(path, "supervised_final")
         self.tokenizer = AutoTokenizer.from_pretrained(self.sup_dir, use_fast=True)
-        self.model = AutoModelForSequenceClassification.from_pretrained(self.sup_dir)
-        self.model.to(self.device)
+        self.model = AutoModelForSequenceClassification.from_pretrained(self.sup_dir).to(self.device)
         self.model.eval()
 
-        # id2label is often included in config; keep your json too as ground truth
         id2label_path = os.path.join(self.sup_dir, "id2label.json")
         if os.path.exists(id2label_path):
             with open(id2label_path, "r", encoding="utf-8") as f:
                 self.id2label = {int(k): v for k, v in json.load(f).items()}
         else:
-            # fallback to config mapping if json missing
             self.id2label = {int(k): v for k, v in self.model.config.id2label.items()}
 
-        # ---- 2) Semantic retrieval (FAISS) ----
-        self.faiss_index_path = os.path.join(path, "faiss_labels.index")
-        self.label_table_path = os.path.join(path, "label_table.csv")
-
-        if os.path.exists(self.faiss_index_path) and os.path.exists(self.label_table_path):
-            self.index = faiss.read_index(self.faiss_index_path)
-            self.label_table = pd.read_csv(self.label_table_path)
-
-            # Expectation: label_table has at least one column that maps FAISS row -> label
-            # Common patterns: ["label_id"], ["unspsc"], ["code"], etc.
-            # We'll build a robust resolver:
-            self.label_cols = [c.lower() for c in self.label_table.columns]
-        else:
-            self.index = None
-            self.label_table = None
-            self.label_cols = []
-
-        # ---- 3) Bi-encoder model for embedding the query ----
-        biencoder_name_path = os.path.join(path, "biencoder_model_name.txt")
-        self.biencoder_name = None
+        # ---- 2) Bi-encoder ----
         self.biencoder = None
+        self.biencoder_name = None
+        biencoder_name_path = os.path.join(path, "biencoder_model_name.txt")
         if os.path.exists(biencoder_name_path):
             with open(biencoder_name_path, "r", encoding="utf-8") as f:
                 self.biencoder_name = f.read().strip()
             if self.biencoder_name:
                 self.biencoder = SentenceTransformer(self.biencoder_name, device=self.device)
 
-        # ---- 4) Optional cross-encoder reranker ----
-        rerank_name_path = os.path.join(path, "rerank_model_name.txt")
-        self.rerank_name = None
+        # ---- 3) Semantic #1: label FAISS ----
+        self.seg_index = None
+        self.seg_table = None
+        self.seg_text_col = None
+
+        seg_index_path = os.path.join(path, "faiss_labels.index")
+        seg_table_path = os.path.join(path, "label_table.csv")
+        if os.path.exists(seg_index_path) and os.path.exists(seg_table_path):
+            self.seg_index = faiss.read_index(seg_index_path)
+            self.seg_table = pd.read_csv(seg_table_path)
+            self.seg_text_col = "label_text" if "label_text" in self.seg_table.columns else None
+
+        # ---- 4) Semantic #2: catalog FAISS ----
+        self.cat_index = None
+        self.cat_table = None
+        self.cat_text_col = None
+
+        cat_index_path = os.path.join(path, "faiss_catalog_articles.index")
+        cat_table_path = os.path.join(path, "catalog_articles_table.csv")
+        if os.path.exists(cat_index_path) and os.path.exists(cat_table_path):
+            self.cat_index = faiss.read_index(cat_index_path)
+            self.cat_table = pd.read_csv(cat_table_path)
+            self.cat_text_col = "doc_text" if "doc_text" in self.cat_table.columns else None
+
+        # ---- 5) Optional reranker ----
         self.reranker = None
+        self.rerank_name = None
+        rerank_name_path = os.path.join(path, "rerank_model_name.txt")
         if os.path.exists(rerank_name_path):
             with open(rerank_name_path, "r", encoding="utf-8") as f:
                 self.rerank_name = f.read().strip()
             if self.rerank_name:
                 self.reranker = CrossEncoder(self.rerank_name, device=self.device)
 
-    # ---------- Helpers ----------
-    def _label_from_table_row(self, row: pd.Series) -> str:
-    # FAISS row i corresponds to one label/code
-        return str(row["s4_code"])
+    # -----------------------------
+    # Query formatting
+    # -----------------------------
+    def _make_query_text(self, text: str, uom: Optional[str] = None) -> str:
+        # classifier-friendly format (keep UOM)
+        t = (text or "").strip()
+        u = (uom or "UN").strip().upper()
+        return f"[UOM={u}] {t}".strip()
 
+    def _make_semantic_query(self, raw_text: str) -> str:
+        """
+        Semantic input cleaner (both semantics):
+        - keep only what is after first '|'
+        - replace newlines by spaces
+        """
+        return _strip_after_first_bar(raw_text)
 
-    def _supervised_topk(self, text: str, top_k: int) -> List[Dict[str, Any]]:
+    # -----------------------------
+    # Supervised topK
+    # -----------------------------
+    def _supervised_topk(self, text: str, top_k: int = 20) -> List[Dict[str, Any]]:
         inputs = self.tokenizer(
             text,
             truncation=True,
@@ -101,127 +248,275 @@ class EndpointHandler:
         with torch.no_grad():
             logits = self.model(**inputs).logits[0].detach().cpu().numpy()
 
-        probs = softmax(logits)
+        probs = softmax_1d(logits)
         idx = np.argsort(-probs)[:top_k]
 
         out = []
         for i in idx:
-            out.append(
-                {
-                    "label_id": int(i),
-                    "label": self.id2label.get(int(i), str(i)),
-                    "score": float(probs[i]),
-                }
-            )
+            s4 = self.id2label.get(int(i), str(i))
+            out.append({"label_id": int(i), "s4_code": str(s4), "score": float(probs[i])})
         return out
 
-    def _semantic_topk(self, text: str, top_k: int) -> List[Dict[str, Any]]:
-        if self.index is None or self.label_table is None or self.biencoder is None:
+    # -----------------------------
+    # Semantic retrieval
+    # -----------------------------
+    def _semantic_topk(
+        self,
+        index: Optional[faiss.Index],
+        table: Optional[pd.DataFrame],
+        text_col: Optional[str],
+        raw_query: str,
+        top_k: int = 50,
+        source: str = "semantic",
+        clean_label_text: bool = False,
+    ) -> List[Dict[str, Any]]:
+        if index is None or table is None or self.biencoder is None:
+            return []
+        if "s4_code" not in table.columns:
             return []
 
-        emb = self.biencoder.encode([text], normalize_embeddings=True)
+        #  Clean semantic query for both semantics
+        query = self._make_semantic_query(raw_query)
+        if not query:
+            # fallback to raw query without bar stripping if empty
+            query = re.sub(r"\s+", " ", (raw_query or "").replace("\n", " ").replace("\r", " ")).strip()
+
+        emb = self.biencoder.encode([query], normalize_embeddings=True)
         emb = np.asarray(emb, dtype=np.float32)
 
-        scores, idx = self.index.search(emb, top_k)  # scores: (1,k), idx: (1,k)
+        scores, idx = index.search(emb, top_k)
         scores = scores[0].tolist()
         idx = idx[0].tolist()
 
         out = []
         for s, i in zip(scores, idx):
-            if i < 0 or i >= len(self.label_table):
+            if i < 0 or i >= len(table):
                 continue
-            row = self.label_table.iloc[i]
+            row = table.iloc[i]
+
+            txt = None
+            if text_col and text_col in table.columns:
+                txt = str(row[text_col])
+
+            #  Extra cleaning for semantic label candidate texts
+            txt_clean = None
+            if txt is not None:
+                txt_one_line = txt.replace("\r", " ").replace("\n", " ")
+                if source == "semantic_label":
+                    txt_clean = clean_custom_code_name(txt_one_line) if clean_label_text else txt_one_line
+                else:
+                    # catalog: at least strip after bar + newline
+                    txt_clean = _strip_after_first_bar(txt_one_line)
+
             out.append(
                 {
+                    "source": source,
                     "faiss_id": int(i),
                     "s4_code": str(row["s4_code"]),
-                    "label_text": str(row["label_text"]),
+                    "label_text": txt,
+                    "label_text_clean": txt_clean,
                     "score": float(s),
+                    "score_01": norm_cos_to_01(float(s)),
+                    "semantic_query": query,
                 }
             )
-
         return out
 
-    def _rerank(self, query: str, candidates: List[Dict[str, Any]], top_k: int) -> List[Dict[str, Any]]:
-        """
-        Rerank semantic candidates using cross-encoder if available.
-        Assumes candidate["label"] is a text representation.
-        """
-        if self.reranker is None or not candidates:
-            return candidates[:top_k]
+    # -----------------------------
+    # Fusion + optional rerank
+    # -----------------------------
+    def _fuse_and_choose(
+        self,
+        raw_query: str,
+        sup: List[Dict[str, Any]],
+        seg: List[Dict[str, Any]],
+        cat: List[Dict[str, Any]],
+        use_rerank: bool,
+        weights: Tuple[float, float, float, float] = (1.0, 0.6, 0.8, 0.5),
+        rerank_topn: int = 40,
+        debug_topn: int = 10,
+    ) -> Dict[str, Any]:
+        w_sup, w_seg, w_cat, w_rerank = weights
+
+        # maps
+        sup_map = {x["s4_code"]: float(x["score"]) for x in sup}
+        seg_map = defaultdict(float)
+        cat_map = defaultdict(float)
+
+        # candidate texts for rerank (prefer cleaned)
+        text_map = {}
+
+        for x in seg:
+            code = x["s4_code"]
+            seg_map[code] = max(seg_map[code], float(x.get("score_01", 0.0)))
+            txt = x.get("label_text_clean") or x.get("label_text") or code
+            text_map[code] = txt
+
+        for x in cat:
+            code = x["s4_code"]
+            cat_map[code] = max(cat_map[code], float(x.get("score_01", 0.0)))
+            txt = x.get("label_text_clean") or x.get("label_text") or code
+            # keep existing if seg provided (often better), else catalog
+            text_map.setdefault(code, txt)
+
+        candidates = set(sup_map.keys()) | set(seg_map.keys()) | set(cat_map.keys())
+        if not candidates:
+            return {"prediction": None, "method": "no_candidates", "fused_topk": []}
+
+        fused = []
+        for code in candidates:
+            s_sup = sup_map.get(code, 0.0)
+            s_seg = seg_map.get(code, 0.0)
+            s_cat = cat_map.get(code, 0.0)
+
+            base = (w_sup * s_sup) + (w_seg * s_seg) + (w_cat * s_cat)
+            fused.append(
+                {
+                    "s4_code": code,
+                    "score": float(base),
+                    "components": {"sup": float(s_sup), "seg": float(s_seg), "cat": float(s_cat)},
+                    "text_for_rerank": text_map.get(code, code),
+                }
+            )
+
+        fused.sort(key=lambda x: x["score"], reverse=True)
 
-        pairs = [[query, c["label_text"]] for c in candidates]
-        scores = self.reranker.predict(pairs).tolist()
+        #  semantic query cleaning also for rerank query side
+        rerank_query = self._make_semantic_query(raw_query) or raw_query.replace("\n", " ").replace("\r", " ")
+        rerank_query = re.sub(r"\s+", " ", rerank_query).strip()
 
-        merged = []
-        for c, s in zip(candidates, scores):
-            cc = dict(c)
-            cc["rerank_score"] = float(s)
-            merged.append(cc)
+        if use_rerank and self.reranker is not None and fused:
+            topn = fused[: min(rerank_topn, len(fused))]
+            pairs = [[rerank_query, c["text_for_rerank"]] for c in topn]
+            rr_raw = self.reranker.predict(pairs).tolist()
+            rr_01 = minmax_norm([float(v) for v in rr_raw])
 
-        merged.sort(key=lambda x: x.get("rerank_score", -1e9), reverse=True)
-        return merged[:top_k]
+            for c, rraw, r01 in zip(topn, rr_raw, rr_01):
+                c["rerank_raw"] = float(rraw)
+                c["rerank_01"] = float(r01)
+                c["score"] = float(c["score"] + w_rerank * float(r01))
 
-    # ---------- Main entrypoint ----------
+            fused[: len(topn)] = topn
+            fused.sort(key=lambda x: x["score"], reverse=True)
+
+        return {
+            "prediction": fused[0],
+            "method": "fused_score",
+            "fused_topk": fused[:debug_topn],
+        }
+
+    # -----------------------------
+    # HF entrypoint
+    # -----------------------------
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Expected input payload examples:
-        {"inputs": "desc...", "top_k": 20, "threshold": 0.80}
-        {"inputs": ["desc1", "desc2"], "top_k": 20}
-        """
         inputs = data.get("inputs")
         if inputs is None:
             return {"error": "Missing 'inputs'."}
 
         top_k = int(data.get("top_k", 20))
-        threshold = float(data.get("threshold", 0.80))  # you asked 80% as prod-ready gate
+        sem_top_k = int(data.get("sem_top_k", 50))
+        threshold = float(data.get("threshold", 0.80))  # supervised gate
         use_rerank = bool(data.get("use_rerank", True))
 
+        # optional global uom
+        global_uom = data.get("uom") or data.get("CODE_UNITE") or data.get("code_unite")
+
         if isinstance(inputs, str):
-            return self._predict_one(inputs, top_k=top_k, threshold=threshold, use_rerank=use_rerank)
-        elif isinstance(inputs, list):
+            return self._predict_one(inputs, uom=global_uom, top_k=top_k, sem_top_k=sem_top_k,
+                                     threshold=threshold, use_rerank=use_rerank)
+
+        if isinstance(inputs, dict):
+            text = inputs.get("text") or inputs.get("description") or inputs.get("inputs") or ""
+            uom = inputs.get("uom") or inputs.get("CODE_UNITE") or inputs.get("code_unite") or global_uom
+            return self._predict_one(text, uom=uom, top_k=top_k, sem_top_k=sem_top_k,
+                                     threshold=threshold, use_rerank=use_rerank)
+
+        if isinstance(inputs, list):
+            results = []
+            for x in inputs:
+                if isinstance(x, str):
+                    results.append(self._predict_one(x, uom=global_uom, top_k=top_k, sem_top_k=sem_top_k,
+                                                     threshold=threshold, use_rerank=use_rerank))
+                elif isinstance(x, dict):
+                    text = x.get("text") or x.get("description") or x.get("inputs") or ""
+                    uom = x.get("uom") or x.get("CODE_UNITE") or x.get("code_unite") or global_uom
+                    results.append(self._predict_one(text, uom=uom, top_k=top_k, sem_top_k=sem_top_k,
+                                                     threshold=threshold, use_rerank=use_rerank))
+                else:
+                    results.append({"error": "Each item in inputs list must be str or dict."})
+            return {"results": results}
+
+        return {"error": "'inputs' must be a string, dict, or a list of strings/dicts."}
+
+    def _predict_one(
+        self,
+        text: str,
+        uom: Optional[str],
+        top_k: int,
+        sem_top_k: int,
+        threshold: float,
+        use_rerank: bool,
+    ) -> Dict[str, Any]:
+        # classifier query keeps UOM (good for supervised)
+        clf_query = self._make_query_text(text, uom=uom)
+
+        # supervised
+        sup = self._supervised_topk(clf_query, top_k=top_k)
+
+        # fast return if high confidence
+        if sup and sup[0]["score"] >= threshold:
             return {
-                "results": [
-                    self._predict_one(x, top_k=top_k, threshold=threshold, use_rerank=use_rerank)
-                    for x in inputs
-                ]
+                "input": text,
+                "uom": uom or "UN",
+                "prediction": {"s4_code": sup[0]["s4_code"], "score": float(sup[0]["score"]), "components": {"sup": float(sup[0]["score"]), "seg": 0.0, "cat": 0.0}},
+                "method": "supervised_threshold",
+                "supervised_topk": sup[:10],
+                "semantic_label_topk": [],
+                "semantic_catalog_topk": [],
+                "fused_topk": [],
             }
-        else:
-            return {"error": "'inputs' must be a string or a list of strings."}
-
-    def _predict_one(self, text: str, top_k: int, threshold: float, use_rerank: bool) -> Dict[str, Any]:
-        sup = self._supervised_topk(text, top_k=top_k)
-        sem = self._semantic_topk(text, top_k=top_k)
 
-        if use_rerank and sem:
-            sem = self._rerank(text, sem, top_k=top_k)
-
-        sup_labels = [x["label"] for x in sup]
-        sem_set = set(x["s4_code"] for x in sem)
-
-        # 1) If supervised top1 already high confidence -> return it directly
-        if sup and sup[0]["score"] >= threshold:
-            chosen = sup[0]
-            chosen_method = "supervised_threshold"
-        else:
-            # 2) Intersection rule: pick first supervised label that appears in semantic topK
-            chosen = None
-            for cand in sup:
-                if cand["label"] in sem_set:
-                    chosen = cand
-                    break
-
-            if chosen is not None:
-                chosen_method = "intersection"
-            else:
-                # 3) Fallback: if no intersection, keep supervised top1 (or semantic top1 if you prefer)
-                chosen = sup[0] if sup else ({"label": sem[0]["label"], "score": sem[0]["score"]} if sem else None)
-                chosen_method = "supervised_no_intersection" if sup else "semantic_only"
+        
+        raw_sem_query = text if isinstance(text, str) else str(text)
+
+        seg = self._semantic_topk(
+            index=self.seg_index,
+            table=self.seg_table,
+            text_col=self.seg_text_col,
+            raw_query=raw_sem_query,
+            top_k=sem_top_k,
+            source="semantic_label",
+            clean_label_text=True,   # deeper cleaning for semantic label
+        )
+
+        cat = self._semantic_topk(
+            index=self.cat_index,
+            table=self.cat_table,
+            text_col=self.cat_text_col,
+            raw_query=raw_sem_query,
+            top_k=sem_top_k,
+            source="semantic_catalog",
+            clean_label_text=False,
+        )
+
+        fused = self._fuse_and_choose(
+            raw_query=raw_sem_query,
+            sup=sup,
+            seg=seg,
+            cat=cat,
+            use_rerank=use_rerank,
+            weights=(1.0, 0.6, 0.8, 0.5),
+            rerank_topn=40,
+            debug_topn=10,
+        )
 
         return {
             "input": text,
-            "prediction": chosen,
-            "method": chosen_method,
-            "supervised_topk": sup[:10],   # keep it light; raise if you want
-            "semantic_topk": sem[:10],
+            "uom": uom or "UN",
+            "prediction": fused.get("prediction"),
+            "method": fused.get("method"),
+            "supervised_topk": sup[:10],
+            "semantic_label_topk": seg[:10],
+            "semantic_catalog_topk": cat[:10],
+            "fused_topk": fused.get("fused_topk", []),
         }