Download handler_old.py from I4SUPPLY/COLAS_CLASSIFIER_V2: direct link, hf CLI and curl.
- Browser
- Download file 35.5 kB
-
https://huggingface.co/I4SUPPLY/COLAS_CLASSIFIER_V2/resolve/main/handler_old.py
- Command line
-
hf download hf://I4SUPPLY/COLAS_CLASSIFIER_V2/handler_old.py
-
curl -L -o handler_old.py https://huggingface.co/I4SUPPLY/COLAS_CLASSIFIER_V2/resolve/main/handler_old.py
35.5 kB
| import os | |
| import json | |
| import re | |
| import unicodedata | |
| from typing import Any, Dict, List, Optional, Tuple | |
| from collections import defaultdict | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| import faiss | |
| from sentence_transformers import SentenceTransformer, CrossEncoder | |
| # ============================================================ | |
| # 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 sigmoid(x: float) -> float: | |
| return float(1.0 / (1.0 + np.exp(-x))) | |
| def norm_cos_to_01(cos_sim: float) -> float: | |
| v = (cos_sim + 1.0) / 2.0 | |
| return float(max(0.0, min(1.0, v))) | |
| def faiss_score_to_01(score_raw: float, metric_type: Optional[int]) -> float: | |
| """ | |
| Robust conversion: | |
| - If FAISS is L2 distance: smaller is better -> exp(-dist) | |
| - Else assume similarity-ish: | |
| - if in [-1, 1] => (x+1)/2 | |
| - else sigmoid(score/5) | |
| """ | |
| if metric_type == faiss.METRIC_L2: | |
| return float(np.exp(-float(score_raw))) | |
| s = float(score_raw) | |
| if -1.01 <= s <= 1.01: | |
| return norm_cos_to_01(s) | |
| return float(sigmoid(s / 5.0)) | |
| def parse_bool(v: Any, default: bool = True) -> bool: | |
| if v is None: | |
| return default | |
| if isinstance(v, bool): | |
| return v | |
| if isinstance(v, (int, float)): | |
| return bool(v) | |
| if isinstance(v, str): | |
| return v.strip().lower() in ("1", "true", "yes", "y", "on") | |
| return bool(v) | |
| # ============================================================ | |
| # Cleaning helpers | |
| # ============================================================ | |
| def strip_after_first_bar_keep_newlines(s: str) -> str: | |
| """ | |
| Keep only what is after the first '|', if any, WITHOUT flattening newlines. | |
| Example: | |
| "uom: UN | A\nB" -> "A\nB" | |
| """ | |
| if not s or not isinstance(s, str): | |
| return "" | |
| if "|" in s: | |
| return s.split("|", 1)[1] | |
| return s | |
| def fix_encoding(text: str) -> str: | |
| """ | |
| Best-effort fixer for common mojibake: "associés" -> "associés" | |
| """ | |
| if not text or not isinstance(text, str): | |
| return "" | |
| t = text | |
| try: | |
| t2 = t.encode("latin1", errors="strict").decode("utf-8", errors="strict") | |
| if "Ã" in t or "�" in t: | |
| return t2 | |
| except Exception: | |
| pass | |
| 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 label text (semantic_label candidates) to improve matching. | |
| """ | |
| if not name or not isinstance(name, str): | |
| return "" | |
| name = fix_encoding(name) | |
| name = re.sub(r"\(\d{5,}\)", "", name) # (0006111100) | |
| name = re.sub(r"^NPO\s*-\s*\d+\s*-\s*", "", name, flags=re.IGNORECASE) | |
| name = re.sub(r"^\d+\s*[-:]\s*", "", name) | |
| name = re.sub(r"[()]", " ", name) | |
| name = re.sub(r"[\[\]]", " ", name) | |
| name = re.sub(r"[,;/\\|#@&*=+]", " ", name) | |
| name = re.sub(r"[-–—]", " ", name) | |
| name = re.sub(r"\s+", " ", name).strip() | |
| name = unicodedata.normalize("NFC", name) | |
| return name | |
| # ============================================================ | |
| # GLOBAL semantic query cleaner with newline usefulness check | |
| # ============================================================ | |
| SEM_RE_MULTI_SPACE = re.compile(r"\s+") | |
| SEM_RE_BRACKET_BLOCK = re.compile(r"\[[^\]]{1,60}\]") # remove [649079], [HUT], etc. | |
| # trailing "(...)" that looks like a code/id (allows . and / and - and _) | |
| SEM_RE_TRAILING_PARENS_CODELIKE = re.compile(r"\s*\(\s*[A-Za-z0-9][A-Za-z0-9._/\-]{3,60}\s*\)\s*$") | |
| # common "ref/code/id" patterns | |
| SEM_RE_JUNK_CODES = re.compile(r"(?i)\b(ref|réf|reference|référence|code|article|art|sku|id)\s*[:#]?\s*[A-Za-z0-9][A-Za-z0-9./_-]{2,}\b") | |
| # units / measures that are typically noise in semantic query | |
| SEM_RE_UNITS = re.compile( | |
| r"(?i)\b(" | |
| r"kg|kilogramme?s?|g|gr|t|tonne?s?|" | |
| r"mm|millim[eè]tre?s?|cm|centim[eè]tre?s?|m\b|m[èe]tre?s?|" | |
| r"l\b|litre?s?|ml|cl|" | |
| r"bar|pfa|pn|psi|pa|mpa|" | |
| r"w\b|kw|v\b|a\b|ah|hz|" | |
| r"°c|c°" | |
| r")\b" | |
| ) | |
| # attribute-like lines: "Poids :", "Forme :", "Usage :", etc. | |
| SEM_ATTR_KEYS = [ | |
| "poids", "poid", "dimension", "dimensions", "dim", "longueur", "largeur", "hauteur", | |
| "epaisseur", "épaisseur", "diametre", "diamètre", "ø", "diam", | |
| "référence", "reference", "ref", "réf", "code", "article", "sku", "id", | |
| "forme", "usage", "matiere", "matière", "type", "couleur", "coloris", | |
| "pfa", "pn", "pression", "pressure", "température", "temperature", | |
| "conditionnement", "colisage", "unité", "uom", "unite", | |
| ] | |
| SEM_RE_ATTR_KV = re.compile( | |
| r"(?i)\b(" + "|".join(re.escape(k) for k in SEM_ATTR_KEYS) + r")\b\s*[:=]" | |
| ) | |
| SEM_RE_MANY_KV = re.compile(r".*[:=].*[:=].*") # 2+ key/value in same line | |
| SEM_RE_JUST_REFERENCE = re.compile(r"(?i)^\s*(référence|reference|ref|réf)\b") | |
| def _digit_ratio(s: str) -> float: | |
| digits = sum(ch.isdigit() for ch in s) | |
| letters = sum(ch.isalpha() for ch in s) | |
| denom = max(1, digits + letters) | |
| return digits / denom | |
| def _looks_like_attribute_line(line: str) -> bool: | |
| """ | |
| Global heuristic: decide if a newline continuation line is "attribute/noise". | |
| Examples: | |
| "Poids : 4 050,00 KG" | |
| "Forme : droit - Usage : avant compteur - ... - PFA : 16" | |
| "bar" | |
| "Référence 1410/B" | |
| """ | |
| if not line: | |
| return True | |
| l = line.strip() | |
| low = l.lower() | |
| # very short single-token unit-ish line | |
| if len(l) <= 6 and SEM_RE_UNITS.search(l): | |
| return True | |
| # starts with "Référence"/"Ref" | |
| if SEM_RE_JUST_REFERENCE.match(l): | |
| return True | |
| # contains explicit attribute keys followed by ":" or "=" | |
| if SEM_RE_ATTR_KV.search(l): | |
| return True | |
| # multiple key/value in one line (usually specs) | |
| if SEM_RE_MANY_KV.match(l): | |
| return True | |
| # lots of digits or many units | |
| if _digit_ratio(l) >= 0.22 and SEM_RE_UNITS.search(l): | |
| return True | |
| # many separators typical of spec lists | |
| if l.count(" - ") >= 2 and (":" in l or "=" in l): | |
| return True | |
| # line mostly numeric/symbols | |
| if _digit_ratio(l) >= 0.35: | |
| return True | |
| # contains "pfa" or "pn" words even without ":" | |
| if " pfa" in (" " + low + " ") or " pn" in (" " + low + " "): | |
| return True | |
| return False | |
| def _select_useful_newline_parts(raw: str, max_extra_lines: int = 1) -> str: | |
| """ | |
| Keep first line always. | |
| Keep up to `max_extra_lines` additional lines ONLY if they do NOT look like attributes/noise. | |
| """ | |
| t = (raw or "").replace("\r", "\n") | |
| # strip trailing oui/non artifacts | |
| t = re.sub(r"\n\s*(oui|non)\s*$", "", t, flags=re.IGNORECASE).strip() | |
| lines = [ln.strip() for ln in t.split("\n") if ln.strip()] | |
| if not lines: | |
| return "" | |
| kept = [lines[0]] | |
| extras = 0 | |
| for ln in lines[1:]: | |
| if _looks_like_attribute_line(ln): | |
| continue | |
| # keep only if it brings actual words | |
| if len(re.findall(r"[A-Za-zÀ-ÿ]{3,}", ln)) < 2: | |
| continue | |
| kept.append(ln) | |
| extras += 1 | |
| if extras >= max_extra_lines: | |
| break | |
| return " ".join(kept) | |
| def clean_semantic_query( | |
| raw_text: str, | |
| *, | |
| drop_leading_long_number: bool = False, | |
| max_extra_lines: int = 1, | |
| ) -> str: | |
| """ | |
| Global semantic cleaner: | |
| - keep text after first '|' | |
| - newline usefulness filter (drop weight/spec/ref lines) | |
| - remove [...] blocks | |
| - remove trailing (CODE) blocks at end (repeat) | |
| - remove explicit REF/CODE patterns | |
| - normalize whitespace | |
| """ | |
| if not raw_text or not isinstance(raw_text, str): | |
| return "" | |
| t = fix_encoding(raw_text) | |
| # keep only after first '|', keep newlines for analysis | |
| t = strip_after_first_bar_keep_newlines(t) | |
| # select useful newline parts (this is the "method to check if after \\n is useful") | |
| t = _select_useful_newline_parts(t, max_extra_lines=max_extra_lines) | |
| # remove square bracket blocks (IDs, supplier tags, etc.) | |
| t = SEM_RE_BRACKET_BLOCK.sub(" ", t) | |
| # optionally remove leading long numbers (line number / internal id) | |
| if drop_leading_long_number: | |
| t = re.sub(r"^\d{4,}\s+", "", t.strip()) | |
| # remove trailing "(...)" that looks like code/id; repeat if multiple | |
| prev = None | |
| while prev != t: | |
| prev = t | |
| t = SEM_RE_TRAILING_PARENS_CODELIKE.sub("", t).strip() | |
| # remove obvious ref/code patterns | |
| t = SEM_RE_JUNK_CODES.sub(" ", t) | |
| # normalize spaces | |
| t = SEM_RE_MULTI_SPACE.sub(" ", t).strip() | |
| return t | |
| # ============================================================ | |
| # Handler | |
| # ============================================================ | |
| class EndpointHandler: | |
| """ | |
| HF Inference Endpoint Custom Handler (Dual Semantic + Catalog Dominance + Trash handling) | |
| Artifacts expected in repo root: | |
| - supervised_final/ (HF transformers classifier) | |
| - faiss_labels.index + label_table.csv | |
| - faiss_catalog_articles.index + catalog_articles_table.csv | |
| (fallback names also supported: faiss_catalog.index + catalog_table.csv) | |
| - biencoder_model_name.txt | |
| - optional: rerank_model_name.txt | |
| """ | |
| def __init__(self, path: str = ""): | |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # ---- 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).to(self.device) | |
| self.model.eval() | |
| 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: | |
| self.id2label = {int(k): v for k, v in self.model.config.id2label.items()} | |
| # ---- 2) Bi-encoder ---- | |
| self.biencoder = 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: | |
| biencoder_name = f.read().strip() | |
| if biencoder_name: | |
| self.biencoder = SentenceTransformer(biencoder_name, device=self.device) | |
| # ---- 3) Semantic #1: labels FAISS ---- | |
| self.seg_index = None | |
| self.seg_metric = 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_metric = getattr(self.seg_index, "metric_type", None) | |
| self.seg_table = pd.read_csv(seg_table_path) | |
| for c in ["label_text", "doc_text", "text", "name", "description"]: | |
| if c in self.seg_table.columns: | |
| self.seg_text_col = c | |
| break | |
| # ---- 4) Semantic #2: catalog FAISS ---- | |
| self.cat_index = None | |
| self.cat_metric = None | |
| self.cat_table = None | |
| self.cat_text_col = None | |
| cat_index_candidates = [ | |
| os.path.join(path, "faiss_catalog_articles.index"), | |
| os.path.join(path, "faiss_catalog.index"), | |
| ] | |
| cat_table_candidates = [ | |
| os.path.join(path, "catalog_articles_table.csv"), | |
| os.path.join(path, "catalog_table.csv"), | |
| ] | |
| cat_index_path = next((p for p in cat_index_candidates if os.path.exists(p)), None) | |
| cat_table_path = next((p for p in cat_table_candidates if os.path.exists(p)), None) | |
| if cat_index_path and cat_table_path: | |
| self.cat_index = faiss.read_index(cat_index_path) | |
| self.cat_metric = getattr(self.cat_index, "metric_type", None) | |
| self.cat_table = pd.read_csv(cat_table_path) | |
| for c in ["doc_text", "item_text", "label_text", "text", "description", "name", "title"]: | |
| if c in self.cat_table.columns: | |
| self.cat_text_col = c | |
| break | |
| # ---- 5) Optional reranker ---- | |
| self.reranker = 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: | |
| rerank_name = f.read().strip() | |
| if rerank_name: | |
| self.reranker = CrossEncoder( | |
| rerank_name, | |
| max_length=512, | |
| device=self.device, | |
| ) | |
| # ----------------------------- | |
| # Input normalization / query formatting | |
| # ----------------------------- | |
| def _normalize_incoming(self, raw_text: str, uom: Optional[str]) -> Tuple[str, str]: | |
| """ | |
| Normalize user input: | |
| - If input is like "uom: UN | <desc>\nnon" extract UOM and desc. | |
| - Strip trailing "oui/non" artifacts. | |
| Returns (desc, final_uom) | |
| """ | |
| t = (raw_text or "").strip() | |
| m = re.match(r"^\s*uom\s*:\s*([A-Za-z0-9]+)\s*\|\s*(.*)$", t, flags=re.IGNORECASE | re.DOTALL) | |
| if m: | |
| uom = uom or m.group(1).strip().upper() | |
| t = m.group(2).strip() | |
| t = re.sub(r"\n\s*(oui|non)\s*$", "", t, flags=re.IGNORECASE).strip() | |
| t = t.replace("\r", "\n") | |
| t = re.sub(r"[ \t]+", " ", t).strip() | |
| final_uom = (uom or "UN").strip().upper() | |
| return t, final_uom | |
| def _make_query_text(self, desc: str, uom: str) -> str: | |
| return f"[UOM={uom}] {(desc or '').strip()}".strip() | |
| def _make_semantic_query(self, raw_text: str, *, drop_leading_long_number: bool, max_extra_lines: int) -> str: | |
| return clean_semantic_query( | |
| raw_text, | |
| drop_leading_long_number=drop_leading_long_number, | |
| max_extra_lines=max_extra_lines, | |
| ) | |
| # ----------------------------- | |
| # Supervised topK | |
| # ----------------------------- | |
| def _supervised_topk(self, text: str, top_k: int = 20) -> List[Dict[str, Any]]: | |
| inputs = self.tokenizer( | |
| text, | |
| truncation=True, | |
| padding=True, | |
| max_length=256, | |
| return_tensors="pt", | |
| ).to(self.device) | |
| with torch.no_grad(): | |
| logits = self.model(**inputs).logits[0].detach().cpu().numpy() | |
| probs = softmax_1d(logits) | |
| idx = np.argsort(-probs)[:top_k] | |
| out = [] | |
| for i in idx: | |
| s4 = self.id2label.get(int(i), str(i)) | |
| out.append({"label_id": int(i), "s4_code": str(s4), "score": float(probs[i])}) | |
| return out | |
| # ----------------------------- | |
| # Semantic retrieval (generic) | |
| # ----------------------------- | |
| def _semantic_topk( | |
| self, | |
| index: Optional[faiss.Index], | |
| metric_type: Optional[int], | |
| table: Optional[pd.DataFrame], | |
| text_col: Optional[str], | |
| raw_query: str, | |
| *, | |
| top_k: int, | |
| source: str, | |
| semantic_drop_leading_long_number: bool, | |
| semantic_max_extra_lines: int, | |
| clean_label_candidates: bool, | |
| clean_catalog_candidates: bool, | |
| ) -> 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 [] | |
| query = self._make_semantic_query( | |
| raw_query, | |
| drop_leading_long_number=semantic_drop_leading_long_number, | |
| max_extra_lines=semantic_max_extra_lines, | |
| ) | |
| if not query: | |
| return [] | |
| emb = self.biencoder.encode([query], normalize_embeddings=True) | |
| emb = np.asarray(emb, dtype=np.float32) | |
| 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(table): | |
| continue | |
| row = table.iloc[i] | |
| s4 = str(row.get("s4_code", "")).strip() | |
| if not s4: | |
| continue | |
| txt = "" | |
| if text_col and text_col in table.columns: | |
| txt = str(row.get(text_col, "") or "") | |
| txt_one_line = txt.replace("\r", " ").replace("\n", " ").strip() | |
| if source == "semantic_label": | |
| txt_clean = clean_custom_code_name(txt_one_line) if clean_label_candidates else txt_one_line | |
| else: | |
| # catalog candidates: clean similarly (helps rerank + map) | |
| if clean_catalog_candidates: | |
| txt_clean = clean_semantic_query(txt_one_line, drop_leading_long_number=False, max_extra_lines=0) or txt_one_line | |
| else: | |
| txt_clean = txt_one_line | |
| out.append( | |
| { | |
| "source": source, | |
| "faiss_id": int(i), | |
| "s4_code": s4, | |
| "label_text": txt_one_line, | |
| "label_text_clean": txt_clean, | |
| "score_raw": float(s), | |
| "score_01": float(faiss_score_to_01(float(s), metric_type)), | |
| "semantic_query": query, | |
| } | |
| ) | |
| out.sort(key=lambda x: x.get("score_01", 0.0), reverse=True) | |
| return out | |
| # ----------------------------- | |
| # Fusion + optional rerank (RRF rank-based) | |
| # ----------------------------- | |
| 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, | |
| trash_labels: set, | |
| trash_penalty_rrf: float, | |
| weights: Tuple[float, float, float, float] = (1.0, 0.6, 0.8, 0.5), | |
| rerank_topn: int = 40, | |
| debug_topn: int = 10, | |
| rrf_k: int = 60, | |
| semantic_drop_leading_long_number: bool = False, | |
| semantic_max_extra_lines: int = 1, | |
| ) -> Dict[str, Any]: | |
| w_sup, w_seg, w_cat, w_rerank = weights | |
| def rrf(rank: Optional[int]) -> float: | |
| if rank is None: | |
| return 0.0 | |
| return 1.0 / (rrf_k + rank) | |
| def build_rrf_sum_map(items: List[Dict[str, Any]], code_key="s4_code"): | |
| score_map = defaultdict(float) | |
| first_rank = {} | |
| for rank, x in enumerate(items, start=1): | |
| code = str(x.get(code_key) or "") | |
| if not code: | |
| continue | |
| score_map[code] += rrf(rank) | |
| first_rank.setdefault(code, rank) | |
| return score_map, first_rank | |
| sup_score_map, sup_first_rank = build_rrf_sum_map(sup) | |
| seg_score_map, seg_first_rank = build_rrf_sum_map(seg) | |
| cat_score_map, cat_first_rank = build_rrf_sum_map(cat) | |
| # candidate texts for rerank (prefer cleaned) | |
| text_map = {} | |
| for x in seg: | |
| code = str(x["s4_code"]) | |
| txt = x.get("label_text_clean") or x.get("label_text") or code | |
| text_map[code] = txt | |
| for x in cat: | |
| code = str(x["s4_code"]) | |
| txt = x.get("label_text_clean") or x.get("label_text") or code | |
| text_map.setdefault(code, txt) | |
| candidates = set(sup_score_map.keys()) | set(seg_score_map.keys()) | set(cat_score_map.keys()) | |
| if not candidates: | |
| return {"prediction": None, "method": "no_candidates", "fused_topk": []} | |
| fused = [] | |
| for code in candidates: | |
| score = ( | |
| w_sup * sup_score_map.get(code, 0.0) | |
| + w_seg * seg_score_map.get(code, 0.0) | |
| + w_cat * cat_score_map.get(code, 0.0) | |
| ) | |
| # penalize trash labels (RRF-scale penalty) | |
| if code in trash_labels: | |
| score -= float(trash_penalty_rrf) | |
| fused.append( | |
| { | |
| "s4_code": code, | |
| "score": float(score), | |
| "components": { | |
| "sup_rank": sup_first_rank.get(code), | |
| "seg_rank": seg_first_rank.get(code), | |
| "cat_rank": cat_first_rank.get(code), | |
| }, | |
| "text_for_rerank": text_map.get(code, code), | |
| } | |
| ) | |
| fused.sort(key=lambda x: x["score"], reverse=True) | |
| # optional rerank (rank-based, stable) | |
| rerank_query = self._make_semantic_query( | |
| raw_query, | |
| drop_leading_long_number=semantic_drop_leading_long_number, | |
| max_extra_lines=semantic_max_extra_lines, | |
| ) | |
| if not rerank_query: | |
| rerank_query = re.sub(r"\s+", " ", (raw_query or "").replace("\r", " ").replace("\n", " ")).strip() | |
| 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() | |
| order = np.argsort(-np.array(rr_raw)) | |
| rr_rank_map = {int(idx): r for r, idx in enumerate(order, start=1)} | |
| for i, c in enumerate(topn): | |
| c["rerank_raw"] = float(rr_raw[i]) | |
| c["rerank_rank"] = rr_rank_map[i] | |
| c["score"] = float(c["score"] + w_rerank * (1.0 / (60 + c["rerank_rank"]))) | |
| fused[: len(topn)] = topn | |
| fused.sort(key=lambda x: x["score"], reverse=True) | |
| return {"prediction": fused[0], "method": "fused_rrf_rank", "fused_topk": fused[:debug_topn]} | |
| # ----------------------------- | |
| # HF entrypoint | |
| # ----------------------------- | |
| def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: | |
| inputs = data.get("inputs") | |
| if inputs is None: | |
| return {"error": "Missing 'inputs'."} | |
| top_k = int(data.get("top_k", 20)) | |
| sem_top_k = int(data.get("sem_top_k", 50)) | |
| # No supervised early return; threshold becomes a gate | |
| sup_min_for_fusion = float(data.get("sup_min_for_fusion", data.get("threshold", 0.80))) | |
| sup_margin_min = float(data.get("sup_margin_min", 0.03)) | |
| sup_fusion_topn = int(data.get("sup_fusion_topn", 5)) | |
| # Semantic cleaning knobs (NEW) | |
| semantic_max_extra_lines = int(data.get("semantic_max_extra_lines", 1)) # keep at most 1 non-attribute line after \n | |
| semantic_drop_leading_long_number = parse_bool(data.get("semantic_drop_leading_long_number", False), default=False) | |
| # Catalog dominance knobs | |
| cat_dom_min = float(data.get("cat_dom_min", 0.83)) | |
| cat_dom_gap = float(data.get("cat_dom_gap", 0.02)) | |
| cat_dom_rerank_topn = int(data.get("cat_dom_rerank_topn", 10)) | |
| use_rerank = parse_bool(data.get("use_rerank", True), default=True) | |
| # Trash labels handling | |
| trash_labels = set(data.get("trash_labels", ["AP0009"])) | |
| # RRF score at rank1 is ~ 1/(60+1)=0.016. Typical penalty: 0.01-0.03. | |
| trash_penalty_rrf = float(data.get("trash_penalty_rrf", 0.02)) | |
| trash_sup_min_for_fusion = float(data.get("trash_sup_min_for_fusion", 0.98)) | |
| trash_sup_margin_min = float(data.get("trash_sup_margin_min", 0.15)) | |
| # optional global uom | |
| global_uom = data.get("uom") or data.get("CODE_UNITE") or data.get("code_unite") | |
| def run_one(x: Any) -> Dict[str, Any]: | |
| if isinstance(x, str): | |
| return self._predict_one( | |
| text=x, | |
| uom=global_uom, | |
| top_k=top_k, | |
| sem_top_k=sem_top_k, | |
| use_rerank=use_rerank, | |
| sup_min_for_fusion=sup_min_for_fusion, | |
| sup_margin_min=sup_margin_min, | |
| sup_fusion_topn=sup_fusion_topn, | |
| cat_dom_min=cat_dom_min, | |
| cat_dom_gap=cat_dom_gap, | |
| cat_dom_rerank_topn=cat_dom_rerank_topn, | |
| trash_labels=trash_labels, | |
| trash_penalty_rrf=trash_penalty_rrf, | |
| trash_sup_min_for_fusion=trash_sup_min_for_fusion, | |
| trash_sup_margin_min=trash_sup_margin_min, | |
| semantic_max_extra_lines=semantic_max_extra_lines, | |
| semantic_drop_leading_long_number=semantic_drop_leading_long_number, | |
| ) | |
| if 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 | |
| return self._predict_one( | |
| text=text, | |
| uom=uom, | |
| top_k=top_k, | |
| sem_top_k=sem_top_k, | |
| use_rerank=use_rerank, | |
| sup_min_for_fusion=sup_min_for_fusion, | |
| sup_margin_min=sup_margin_min, | |
| sup_fusion_topn=sup_fusion_topn, | |
| cat_dom_min=cat_dom_min, | |
| cat_dom_gap=cat_dom_gap, | |
| cat_dom_rerank_topn=cat_dom_rerank_topn, | |
| trash_labels=trash_labels, | |
| trash_penalty_rrf=trash_penalty_rrf, | |
| trash_sup_min_for_fusion=trash_sup_min_for_fusion, | |
| trash_sup_margin_min=trash_sup_margin_min, | |
| semantic_max_extra_lines=semantic_max_extra_lines, | |
| semantic_drop_leading_long_number=semantic_drop_leading_long_number, | |
| ) | |
| return {"error": "Each item in inputs must be str or dict."} | |
| if isinstance(inputs, (str, dict)): | |
| return run_one(inputs) | |
| if isinstance(inputs, list): | |
| return {"results": [run_one(x) for x in inputs]} | |
| 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, | |
| use_rerank: bool, | |
| sup_min_for_fusion: float, | |
| sup_margin_min: float, | |
| sup_fusion_topn: int, | |
| cat_dom_min: float, | |
| cat_dom_gap: float, | |
| cat_dom_rerank_topn: int, | |
| trash_labels: set, | |
| trash_penalty_rrf: float, | |
| trash_sup_min_for_fusion: float, | |
| trash_sup_margin_min: float, | |
| semantic_max_extra_lines: int, | |
| semantic_drop_leading_long_number: bool, | |
| ) -> Dict[str, Any]: | |
| desc, final_uom = self._normalize_incoming(text, uom=uom) | |
| clf_query = self._make_query_text(desc, uom=final_uom) | |
| # supervised | |
| sup_all = self._supervised_topk(clf_query, top_k=top_k) | |
| bert_top1 = float(sup_all[0]["score"]) if sup_all else 0.0 | |
| bert_top2 = float(sup_all[1]["score"]) if len(sup_all) > 1 else 0.0 | |
| bert_margin = bert_top1 - bert_top2 | |
| bert_top1_code = str(sup_all[0]["s4_code"]) if sup_all else "" | |
| # gate supervised participation (trash labels require stricter gate) | |
| if bert_top1_code in trash_labels: | |
| bert_used_in_fusion = (bert_top1 >= trash_sup_min_for_fusion) and (bert_margin >= trash_sup_margin_min) | |
| else: | |
| bert_used_in_fusion = (bert_top1 >= sup_min_for_fusion) and (bert_margin >= sup_margin_min) | |
| sup_for_fusion = sup_all[: max(0, int(sup_fusion_topn))] if bert_used_in_fusion else [] | |
| raw_sem_query = text if isinstance(text, str) else str(text) | |
| # semantic label | |
| seg = self._semantic_topk( | |
| index=self.seg_index, | |
| metric_type=self.seg_metric, | |
| table=self.seg_table, | |
| text_col=self.seg_text_col, | |
| raw_query=raw_sem_query, | |
| top_k=sem_top_k, | |
| source="semantic_label", | |
| semantic_drop_leading_long_number=semantic_drop_leading_long_number, | |
| semantic_max_extra_lines=semantic_max_extra_lines, | |
| clean_label_candidates=True, | |
| clean_catalog_candidates=False, | |
| ) | |
| # semantic catalog | |
| cat = self._semantic_topk( | |
| index=self.cat_index, | |
| metric_type=self.cat_metric, | |
| table=self.cat_table, | |
| text_col=self.cat_text_col, | |
| raw_query=raw_sem_query, | |
| top_k=sem_top_k, | |
| source="semantic_catalog", | |
| semantic_drop_leading_long_number=semantic_drop_leading_long_number, | |
| semantic_max_extra_lines=semantic_max_extra_lines, | |
| clean_label_candidates=False, | |
| clean_catalog_candidates=True, | |
| ) | |
| # ----------------------------- | |
| # Catalog dominance override (kept + trash-safe) | |
| # ----------------------------- | |
| if cat: | |
| top1 = cat[0] | |
| top1_code = str(top1["s4_code"]) | |
| top1_s = float(top1.get("score_01", 0.0)) | |
| top2_s = 0.0 | |
| for x in cat[1:]: | |
| if str(x["s4_code"]) != top1_code: | |
| top2_s = float(x.get("score_01", 0.0)) | |
| break | |
| if (top1_code not in trash_labels) and (top1_s >= cat_dom_min) and ((top1_s - top2_s) >= cat_dom_gap): | |
| best = top1 | |
| if use_rerank and self.reranker is not None: | |
| rerank_query = self._make_semantic_query( | |
| raw_sem_query, | |
| drop_leading_long_number=semantic_drop_leading_long_number, | |
| max_extra_lines=semantic_max_extra_lines, | |
| ) | |
| if not rerank_query: | |
| rerank_query = re.sub(r"\s+", " ", raw_sem_query).strip() | |
| unique = [] | |
| seen = set() | |
| for x in cat[: max(1, int(cat_dom_rerank_topn))]: | |
| code = str(x["s4_code"]) | |
| if code in seen: | |
| continue | |
| seen.add(code) | |
| unique.append(x) | |
| pairs = [ | |
| [rerank_query, (u.get("label_text_clean") or u.get("label_text") or str(u["s4_code"]))] | |
| for u in unique | |
| ] | |
| rr_raw = self.reranker.predict(pairs).tolist() | |
| best_idx = int(np.argmax(np.array(rr_raw))) | |
| best = unique[best_idx] | |
| return { | |
| "input": text, | |
| "uom": final_uom, | |
| "prediction": { | |
| "s4_code": str(best["s4_code"]), | |
| "score": float(best.get("score_01", 0.0)), | |
| "components": {"catalog_dom": float(best.get("score_01", 0.0))}, | |
| "text_for_rerank": best.get("label_text_clean") or best.get("label_text"), | |
| "rerank_raw": float(rr_raw[best_idx]), | |
| }, | |
| "method": "catalog_dominance_rerank", | |
| "supervised_topk": sup_all[:10], | |
| "semantic_label_topk": seg[:10], | |
| "semantic_catalog_topk": cat[:10], | |
| "fused_topk": [], | |
| "bert_used_in_fusion": False, | |
| "bert_top1": bert_top1, | |
| "bert_margin": bert_margin, | |
| "semantic_query_used": self._make_semantic_query( | |
| raw_sem_query, | |
| drop_leading_long_number=semantic_drop_leading_long_number, | |
| max_extra_lines=semantic_max_extra_lines, | |
| ), | |
| } | |
| return { | |
| "input": text, | |
| "uom": final_uom, | |
| "prediction": { | |
| "s4_code": top1_code, | |
| "score": top1_s, | |
| "components": {"catalog_dom": top1_s}, | |
| "text_for_rerank": top1.get("label_text_clean") or top1.get("label_text"), | |
| }, | |
| "method": "catalog_dominance", | |
| "supervised_topk": sup_all[:10], | |
| "semantic_label_topk": seg[:10], | |
| "semantic_catalog_topk": cat[:10], | |
| "fused_topk": [], | |
| "bert_used_in_fusion": False, | |
| "bert_top1": bert_top1, | |
| "bert_margin": bert_margin, | |
| "semantic_query_used": self._make_semantic_query( | |
| raw_sem_query, | |
| drop_leading_long_number=semantic_drop_leading_long_number, | |
| max_extra_lines=semantic_max_extra_lines, | |
| ), | |
| } | |
| # ----------------------------- | |
| # Fusion (RRF) with trash penalty | |
| # ----------------------------- | |
| fused = self._fuse_and_choose( | |
| raw_query=raw_sem_query, | |
| sup=sup_for_fusion, | |
| seg=seg, | |
| cat=cat, | |
| use_rerank=use_rerank, | |
| trash_labels=trash_labels, | |
| trash_penalty_rrf=trash_penalty_rrf, | |
| weights=(1.0, 0.6, 0.8, 0.5), | |
| rerank_topn=40, | |
| debug_topn=10, | |
| semantic_drop_leading_long_number=semantic_drop_leading_long_number, | |
| semantic_max_extra_lines=semantic_max_extra_lines, | |
| ) | |
| pred = fused.get("prediction") | |
| # optional trash avoid if close tie | |
| if pred and str(pred.get("s4_code", "")) in trash_labels: | |
| for alt in fused.get("fused_topk", [])[1:]: | |
| if str(alt.get("s4_code", "")) not in trash_labels: | |
| if float(pred.get("score", 0.0)) - float(alt.get("score", 0.0)) <= 0.01: | |
| pred = alt | |
| fused["method"] = "fused_rrf_rank_trash_avoid" | |
| break | |
| return { | |
| "input": text, | |
| "uom": final_uom, | |
| "prediction": pred, | |
| "method": fused.get("method"), | |
| "supervised_topk": sup_all[:10], | |
| "semantic_label_topk": seg[:10], | |
| "semantic_catalog_topk": cat[:10], | |
| "fused_topk": fused.get("fused_topk", []), | |
| "bert_used_in_fusion": bert_used_in_fusion, | |
| "bert_top1": bert_top1, | |
| "bert_margin": bert_margin, | |
| "semantic_query_used": self._make_semantic_query( | |
| raw_sem_query, | |
| drop_leading_long_number=semantic_drop_leading_long_number, | |
| max_extra_lines=semantic_max_extra_lines, | |
| ), | |
| "semantic_max_extra_lines": int(semantic_max_extra_lines), | |
| "semantic_drop_leading_long_number": bool(semantic_drop_leading_long_number), | |
| "trash_labels": sorted(list(trash_labels)), | |
| } |