Update handler.py

handler.py

@@ -335,73 +335,99 @@ class EndpointHandler:
         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,  # bigger = ranks matter a bit less
     ) -> 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)
-
+    
+        def rrf(rank: Optional[int]) -> float:
+            # rank is 1-based. Missing => 0 contribution
+            if rank is None:
+                return 0.0
+            return 1.0 / (rrf_k + rank)
+    
+        def build_rank_map(items: List[Dict[str, Any]], code_key="s4_code") -> Dict[str, int]:
+            # assumes items already ordered best->worst
+            out = {}
+            r = 1
+            for x in items:
+                code = str(x.get(code_key))
+                if code and code not in out:
+                    out[code] = r
+                    r += 1
+            return out
+    
+        # --- build rank maps (order matters) ---
+        sup_rank = build_rank_map(sup)
+    
+        # seg/cat are already in FAISS best->worst order
+        seg_rank = build_rank_map(seg)
+        cat_rank = build_rank_map(cat)
+    
         # 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)))
+            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 = x["s4_code"]
-            cat_map[code] = max(cat_map[code], float(x.get("score_01", 0.0)))
+            code = str(x["s4_code"])
             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())
+    
+        candidates = set(sup_rank.keys()) | set(seg_rank.keys()) | set(cat_rank.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),
-                }
+            # rank-based fusion (scale-free)
+            score = (
+                w_sup * rrf(sup_rank.get(code))
+                + w_seg * rrf(seg_rank.get(code))
+                + w_cat * rrf(cat_rank.get(code))
             )
-
+    
+            fused.append({
+                "s4_code": code,
+                "score": float(score),
+                "components": {
+                    # keep some debug info (ranks + raw supervised prob if exists)
+                    "sup_rank": sup_rank.get(code),
+                    "seg_rank": seg_rank.get(code),
+                    "cat_rank": cat_rank.get(code),
+                },
+                "text_for_rerank": text_map.get(code, code),
+            })
+    
         fused.sort(key=lambda x: x["score"], reverse=True)
-
-        #  semantic query cleaning also for rerank query side
+    
+        # --- optional rerank: use rank, NOT minmax ---
         rerank_query = self._make_semantic_query(raw_query) or raw_query.replace("\n", " ").replace("\r", " ")
         rerank_query = re.sub(r"\s+", " ", rerank_query).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()
-            rr_01 = minmax_norm([float(v) for v in rr_raw])
-
-            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))
-
+    
+            # rank reranker outputs (best = rank 1)
+            order = np.argsort(-np.array(rr_raw))
+            rr_rank_map = {}
+            for r, idx in enumerate(order, start=1):
+                rr_rank_map[idx] = r
+    
+            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 * rrf(c["rerank_rank"]))
+    
             fused[: len(topn)] = topn
             fused.sort(key=lambda x: x["score"], reverse=True)
-
+    
         return {
             "prediction": fused[0],
-            "method": "fused_score",
+            "method": "fused_rrf_rank",
             "fused_topk": fused[:debug_topn],
         }