Spaces:
Configuration error
Configuration error
ML Assistant
feat(upgrades): zerank-2 text reranker tier + Tomoro ColQwen3 visual lane swap + Qwen3-VL-Reranker page reranking (A/B flags)
c15ad20 | """Visual reranker: Qwen3-VL-Reranker-2B (Apache-2.0) scores (query, page | |
| image) pairs and reorders the fused top pages. Query-time only (no rebuild). | |
| Fail-open: on ANY error the input order is returned unchanged.""" | |
| from pathlib import Path | |
| import config | |
| _PROMPT = "Retrieve images or text relevant to the user's query." | |
| _re = None | |
| def get_visual_reranker(): | |
| """Lazy singleton; returns None when loading fails (fail-open upstream).""" | |
| global _re | |
| if _re is False: | |
| return None | |
| if _re is None: | |
| try: | |
| from sentence_transformers import CrossEncoder | |
| _re = CrossEncoder(config.VISUAL_RERANK_MODEL, device=config.DEVICE) | |
| except Exception: | |
| _re = False | |
| return _re or None | |
| def _png(sha, page): | |
| import ingest_lib | |
| return ingest_lib.PAGES_DIR / str(sha) / f"p{int(page):04d}.png" | |
| def rerank_keys(query, keys): | |
| """keys: [(sha, page)] -> reordered list. Entries without an existing | |
| page PNG keep their relative order at the end.""" | |
| if not keys: | |
| return keys | |
| re = get_visual_reranker() | |
| if re is None: | |
| return keys | |
| try: | |
| pairs = [] | |
| scored_idx = [] | |
| rest = [] | |
| for i, (sha, page) in enumerate(keys): | |
| png = None | |
| if sha and page is not None: | |
| try: | |
| png = _png(sha, page) | |
| except Exception: | |
| png = None | |
| if png is not None and Path(png).exists(): | |
| pairs.append((query, {"image": str(png)})) | |
| scored_idx.append(i) | |
| else: | |
| rest.append((i, (sha, page))) | |
| if not pairs: | |
| return keys | |
| try: | |
| scores = re.predict(pairs, prompt=_PROMPT) | |
| except TypeError: | |
| scores = re.predict(pairs) | |
| order = sorted(zip(scored_idx, [float(s) for s in scores]), | |
| key=lambda t: -t[1]) | |
| out = [keys[i] for i, _ in order] | |
| out.extend(k for _, k in rest) | |
| return out | |
| except Exception: | |
| return keys | |