Merge pull request #108 from DataScienceUIBK/demo

fix: correct Diver/Reasonir model registry, remove arena BM25 selecto…

Author: abdoelsayed2016

demo-web/src/app/agent/page.tsx

@@ -126,7 +126,7 @@ export default function AgentPage() {
                                                         strong: ({ node, ...props }) => <strong className="font-bold text-slate-900" {...props} />,
                                                         code: ({ node, inline, ...props }: any) =>
                                                             inline ? (
-                                                                <code className="bg-slate-100 text-pink-600 px-1.5 py-0.5 rounded-md text-[13px] font-mono" {...props} />
+                                                                <code className="bg-slate-100 text-violet-700 px-1.5 py-0.5 rounded-md text-[13px] font-mono font-semibold" {...props} />
                                                             ) : (
                                                                 <div className="my-4 rounded-xl overflow-hidden border border-slate-200 bg-[#0d1117] shadow-sm">
                                                                     <div className="flex items-center px-4 py-2 bg-slate-800 border-b border-slate-700">
@@ -136,8 +136,8 @@ export default function AgentPage() {
                                                                             <div className="w-3 h-3 rounded-full bg-green-500/80"></div>
                                                                         </div>
                                                                     </div>
-                                                                    <pre className="p-4 overflow-x-auto">
-                                                                        <code className="text-[13px] font-mono text-slate-50 leading-relaxed" {...props} />
+                                                                    <pre className="p-4 overflow-x-auto bg-[#0d1117]">
+                                                                        <code className="text-[13px] font-mono text-white leading-relaxed whitespace-pre-wrap" {...props} />
                                                                     </pre>
                                                                 </div>
                                                             ),

demo-web/src/app/arena/page.tsx

@@ -157,11 +157,10 @@ export default function ArenaPage() {
                                 Pipeline A (Baseline)
                             </div>
                             <div className="flex flex-col gap-4">
-                                <Sel value={pipeA.retriever} onChange={() => { }} opts={[{ value: "bm25", label: "BM25 (Sparse Base)" }]} label="Retriever (Fixed by BEIR)" icon={Search} />
                                 <Sel
                                     value={pipeA.method}
                                     onChange={v => setPipeA(p => ({ ...p, method: v, model: v === "none" ? "none" : (RERANKERS_MAP[v as keyof typeof RERANKERS_MAP]?.[0] || "") }))}
-                                    opts={METHODS.map(m => ({ value: m, label: m === "none" ? "None (Base BM25 Only)" : m }))}
+                                    opts={METHODS.map(m => ({ value: m, label: m === "none" ? "None (BM25 baseline only)" : m }))}
                                     label="Reranking Method"
                                     icon={ListTree}
                                 />
@@ -184,11 +183,10 @@ export default function ArenaPage() {
                                 Pipeline B (Challenger)
                             </div>
                             <div className="flex flex-col gap-4">
-                                <Sel value={pipeB.retriever} onChange={() => { }} opts={[{ value: "bm25", label: "BM25 (Sparse Base)" }]} label="Retriever (Fixed by BEIR)" icon={Search} />
                                 <Sel
                                     value={pipeB.method}
                                     onChange={v => setPipeB(p => ({ ...p, method: v, model: v === "none" ? "none" : (RERANKERS_MAP[v as keyof typeof RERANKERS_MAP]?.[0] || "") }))}
-                                    opts={METHODS.map(m => ({ value: m, label: m === "none" ? "None (Base BM25 Only)" : m }))}
+                                    opts={METHODS.map(m => ({ value: m, label: m === "none" ? "None (BM25 baseline only)" : m }))}
                                     label="Reranking Method"
                                     icon={ListTree}
                                 />

demo_server.py

@@ -498,81 +498,60 @@ async def arena_run(req: ArenaRequest):
         eval_docs = random.sample(documents, min(req.n_queries, len(documents)))
 
         def evaluate_pipeline(pipeline_cfg: ArenaPipeline, docs):
-            import copy
+            import copy, math
             docs_copy = copy.deepcopy(docs)
 
-            # Retrieval
-            idx_type = "msmarco" if req.dataset == "msmarco" else "wiki"
-            retriever = get_retriever(pipeline_cfg.retriever, n_docs=req.n_docs, index_type=idx_type)
-            t0 = time.time()
-            ret_results = retriever.retrieve(docs_copy)
-            ret_latency = (time.time() - t0) * 1000 / len(docs_copy)
+            # NOTE: BEIR datasets are already pre-retrieved with BM25 – the downloaded
+            # JSON files contain ranked contexts with `has_answer` set.
+            # Re-calling retriever.retrieve() would reset those contexts and lose the
+            # relevance labels, making all metrics come out as 0.
+            # So we skip re-retrieval; we only rerank if a category is configured.
+            ret_latency = 0.0
+            rr_latency = 0.0
+            ret_results = docs_copy
 
             # Reranking
-            rr_latency = 0
             reranker = get_reranker(pipeline_cfg.rerankerCategory, pipeline_cfg.rerankerModel)
             if reranker:
                 t1 = time.time()
                 ret_results = reranker.rank(ret_results)
-                rr_latency = (time.time() - t1) * 1000 / len(docs_copy)
+                rr_latency = (time.time() - t1) * 1000 / max(1, len(docs_copy))
 
-            # Evaluation - use true TREC evaluation as per BEIR standards
-            metrics = Metrics(ret_results)
+            # Evaluate: pure-Python NDCG@10 and MRR@10 using has_answer flags
             use_rr = reranker is not None
+            mrr_sum = 0.0
+            ndcg_sum = 0.0
 
-            # Formulate the correct qrel name for rankify
-            qrel_name = req.dataset
-            if req.dataset.startswith("beir-"):
-                qrel_name = req.dataset.split("-")[1]
-            elif req.dataset in ["nq-dev", "msmarco", "triviaqa"]:
-                qrel_name = req.dataset
+            for doc in ret_results:
+                contexts = doc.reorder_contexts if (use_rr and getattr(doc, "reorder_contexts", None)) else doc.contexts
+                if not contexts:
+                    continue
 
-            try:
-                trec_metrics = metrics.calculate_trec_metrics(
-                    ndcg_cuts=[10], 
-                    map_cuts=[10], 
-                    mrr_cuts=[10], 
-                    qrel=qrel_name, 
-                    use_reordered=use_rr
-                )
-                ndcg_10 = trec_metrics.get("ndcg@10", 0) * 100
-                mrr_10 = trec_metrics.get("mrr@10", 0) * 100
-            except Exception as e:
-                logger.error(f"TREC Eval Error: {e}")
-                ndcg_10, mrr_10 = 0, 0
-
-            # FALLBACK: If pyserini fails or returns 0.0 (happens on Python 3.13), 
-            # we use a manual calculation based on doc.has_answer.
-            if ndcg_10 == 0 and mrr_10 == 0:
-                logger.warning(f"TREC Eval returned 0.0 for {qrel_name}, using manual fallback.")
-                import math
-                mrr_sum = 0
-                ndcg_sum = 0
-                for doc in ret_results:
-                    contexts = doc.reorder_contexts if (use_rr and doc.reorder_contexts) else doc.contexts
-                    # MRR
-                    found_at = -1
-                    for i, ctx in enumerate(contexts[:10]):
-                        if getattr(ctx, "has_answer", False):
-                            found_at = i + 1
-                            break
-                    if found_at > 0: mrr_sum += 1.0 / found_at
-                    
-                    # NDCG (Binary)
-                    dcg = 0
-                    hits_rels = []
-                    for i, ctx in enumerate(contexts[:10]):
-                        rel = 1 if getattr(ctx, "has_answer", False) else 0
-                        hits_rels.append(rel)
-                        if rel: dcg += 1.0 / math.log2(i + 2)
-                    
-                    hits_rels.sort(reverse=True)
-                    idcg = sum(1.0 / math.log2(i + 2) for i, rel in enumerate(hits_rels) if rel)
-                    if idcg > 0: ndcg_sum += (dcg / idcg)
+                # MRR@10
+                for i, ctx in enumerate(contexts[:10]):
+                    if getattr(ctx, "has_answer", False):
+                        mrr_sum += 1.0 / (i + 1)
+                        break
+                
+                # NDCG@10 (binary relevance)
+                dcg = 0.0
+                rels = []
+                for i, ctx in enumerate(contexts[:10]):
+                    rel = 1 if getattr(ctx, "has_answer", False) else 0
+                    rels.append(rel)
+                    if rel:
+                        dcg += 1.0 / math.log2(i + 2)
 
-                n = len(ret_results)
-                mrr_10 = (mrr_sum / n) * 100 if n > 0 else 0
-                ndcg_10 = (ndcg_sum / n) * 100 if n > 0 else 0
+                rels_sorted = sorted(rels, reverse=True)
+                idcg = sum(r / math.log2(i + 2) for i, r in enumerate(rels_sorted) if r)
+                if idcg > 0:
+                    ndcg_sum += dcg / idcg
+            
+            n = len(ret_results)
+            mrr_10 = (mrr_sum / n) * 100 if n > 0 else 0.0
+            ndcg_10 = (ndcg_sum / n) * 100 if n > 0 else 0.0
+            
+            logger.info(f"Arena eval: n={n} NDCG@10={ndcg_10:.2f}% MRR@10={mrr_10:.2f}%")
 
             return {
                 "mrr_10": mrr_10,

rankify/agent/agent.py

@@ -54,12 +54,39 @@ class RankifyAgent:
 Rankify is a comprehensive Python toolkit for Retrieval, Re-Ranking, and Retrieval-Augmented Generation (RAG).
 
 Your job is to help users select the best models for their use case. You have access to:
-- 10 retrieval methods (BM25, DPR, ANCE, BGE, ColBERT, Contriever, HyDE, Online)
+- **Sparse Retrievers**: BM25 (fast, no GPU, exact match)
+- **Dense Retrievers**: DPR, ANCE, BGE, ColBERT, Contriever
+- **Diver Dense Retrievers** (method="diver-dense"): Many bi-encoder and LLM-based variants selectable via model_id
+- **Reasoning-Augmented Retrievers** (SOTA on BRIGHT benchmark): ReasonIR-8B, ReasonEmbed, BGE-Reasoner-Embed
+- **Online Retriever**: Web search via APIs (real-time data)
+- **HyDE**: Hypothetical Document Embedding for complex queries
 - 23 reranking methods (MonoT5, FlashRank, RankGPT, InRanker, ColBERT, API rerankers, etc.)
 - 7 RAG methods (Basic RAG, Chain-of-Thought, Self-Consistency, ReAct, FiD, etc.)
 
+**Diver Dense Retriever Guide (method="diver-dense"):**
+Valid model_ids (must have corpus_path):
+- `bge` → BAAI/bge-large-en-v1.5
+- `sbert` → sentence-transformers/all-mpnet-base-v2
+- `nomic` → nomic-ai/nomic-embed-text-v1
+- `diver` → AQ-MedAI/Diver-Retriever-4B (flagship diverse evidence model)
+- `inst-l` → hkunlp/instructor-large
+- `inst-xl` → hkunlp/instructor-xl
+- `e5` → intfloat/e5-mistral-7b-instruct (LLM-based)
+- `sf` → Salesforce/SFR-Embedding-Mistral (LLM-based)
+- `rader` → Raderspace/RaDeR_Qwen_25_7B (reasoning-augmented)
+- `grit` → GritLM/GritLM-7B (generative representation)
+- `m2` → togethercomputer/m2-bert-80M-32k-retrieval (long context, 32k)
+- `contriever` → facebook/contriever-msmarco
+
+Example: `Retriever(method="diver-dense", model_id="diver", corpus_path="data/corpus.jsonl", n_docs=10)`
+
+**Reasoning-Augmented Retrievers:**
+- `Retriever(method="reasonir", corpus_path=...)` → reasonir/ReasonIR-8B (SOTA BRIGHT benchmark)
+- `Retriever(method="reason-embed", model_id="qwen3-8b"|"qwen3-4b"|"llama-8b", corpus_path=...)`
+- `Retriever(method="bge-reasoner-embed", corpus_path=...)` → BAAI/bge-reasoner-embed-qwen3-8b
+
 When helping users, consider:
-1. Their task type (QA, search, summarization, conversational)
+1. Their task type (QA, search, summarization, conversational, reasoning-intensive)
 2. Hardware constraints (GPU availability, memory)
 3. Latency requirements
 4. Whether they can use APIs or need local models

rankify/agent/model_registry.py

@@ -187,9 +187,180 @@ def score_for_task(self, task: TaskType) -> float:
         api_provider="serper",
         best_for=["real-time data", "current events", "web search"],
     ),
+
+    # === DIVER DENSE RETRIEVERS ===
+    # method="diver-dense", use model_path as the model_id argument
+    # Valid model_ids: bge, sbert, contriever_st, nomic, diver, inst-l, inst-xl, sf, e5, rader, m2, contriever, grit
+    "diver-bge": ModelMetadata(
+        name="Diver (BGE Large)",
+        method="diver-dense",
+        description="BAAI/bge-large-en-v1.5 via the Diver dense retrieval framework.",
+        speed=Speed.MEDIUM,
+        accuracy=Accuracy.STATE_OF_THE_ART,
+        gpu_required=True,
+        memory_mb=3000,
+        best_for=["semantic search", "high accuracy", "BEIR benchmarks"],
+        model_path="bge",
+    ),
+    "diver-sbert": ModelMetadata(
+        name="Diver (SBERT all-mpnet-base-v2)",
+        method="diver-dense",
+        description="sentence-transformers/all-mpnet-base-v2 via the Diver framework.",
+        speed=Speed.FAST,
+        accuracy=Accuracy.VERY_GOOD,
+        gpu_required=True,
+        memory_mb=1500,
+        best_for=["sentence similarity", "semantic search"],
+        model_path="sbert",
+    ),
+    "diver-nomic": ModelMetadata(
+        name="Diver (Nomic Embed)",
+        method="diver-dense",
+        description="nomic-ai/nomic-embed-text-v1 via the Diver framework.",
+        speed=Speed.FAST,
+        accuracy=Accuracy.VERY_GOOD,
+        gpu_required=True,
+        memory_mb=1500,
+        best_for=["long context", "semantic search", "document retrieval"],
+        model_path="nomic",
+    ),
+    "diver-e5": ModelMetadata(
+        name="Diver (E5-Mistral-7B)",
+        method="diver-dense",
+        description="intfloat/e5-mistral-7b-instruct — instruction-tuned LLM encoder in the Diver framework.",
+        speed=Speed.SLOW,
+        accuracy=Accuracy.STATE_OF_THE_ART,
+        gpu_required=True,
+        memory_mb=16000,
+        best_for=["zero-shot retrieval", "instruction following", "complex queries"],
+        model_path="e5",
+    ),
+    "diver-sf": ModelMetadata(
+        name="Diver (SFR-Embedding-Mistral)",
+        method="diver-dense",
+        description="Salesforce/SFR-Embedding-Mistral — Salesforce Mistral-based bi-encoder in the Diver framework.",
+        speed=Speed.SLOW,
+        accuracy=Accuracy.STATE_OF_THE_ART,
+        gpu_required=True,
+        memory_mb=16000,
+        best_for=["high accuracy retrieval", "complex queries", "BEIR benchmarks"],
+        model_path="sf",
+    ),
+    "diver-rader": ModelMetadata(
+        name="Diver (RaDeR)",
+        method="diver-dense",
+        description="Raderspace/RaDeR_Qwen_25_7B — reasoning-aware dense retriever in the Diver framework.",
+        speed=Speed.SLOW,
+        accuracy=Accuracy.EXCELLENT,
+        gpu_required=True,
+        memory_mb=16000,
+        best_for=["reasoning-intensive queries", "multi-hop QA", "math-related retrieval"],
+        model_path="rader",
+    ),
+    "diver-grit": ModelMetadata(
+        name="Diver (GritLM-7B)",
+        method="diver-dense",
+        description="GritLM/GritLM-7B — generative representation model in the Diver framework.",
+        speed=Speed.VERY_SLOW,
+        accuracy=Accuracy.STATE_OF_THE_ART,
+        gpu_required=True,
+        memory_mb=16000,
+        best_for=["generative retrieval", "LLM-quality embeddings", "long context"],
+        model_path="grit",
+    ),
+    "diver-model": ModelMetadata(
+        name="Diver Retriever-4B",
+        method="diver-dense",
+        description="AQ-MedAI/Diver-Retriever-4B — the flagship Diver diverse-evidence retrieval model.",
+        speed=Speed.SLOW,
+        accuracy=Accuracy.STATE_OF_THE_ART,
+        gpu_required=True,
+        memory_mb=8000,
+        best_for=["diverse evidence retrieval", "BEIR benchmarks", "medical QA"],
+        model_path="diver",
+    ),
+    "diver-inst-l": ModelMetadata(
+        name="Diver (Instructor-Large)",
+        method="diver-dense",
+        description="hkunlp/instructor-large — instruction-following encoder in the Diver framework.",
+        speed=Speed.MEDIUM,
+        accuracy=Accuracy.VERY_GOOD,
+        gpu_required=True,
+        memory_mb=3000,
+        best_for=["instruction following", "domain-specific retrieval"],
+        model_path="inst-l",
+    ),
+    "diver-m2": ModelMetadata(
+        name="Diver (M2-BERT-32K)",
+        method="diver-dense",
+        description="togethercomputer/m2-bert-80M-32k-retrieval — long-context retrieval in the Diver framework.",
+        speed=Speed.MEDIUM,
+        accuracy=Accuracy.VERY_GOOD,
+        gpu_required=True,
+        memory_mb=2000,
+        best_for=["long-context retrieval", "32k sequence length"],
+        model_path="m2",
+    ),
+
+    # === REASONING-AUGMENTED RETRIEVERS ===
+    "reasonir": ModelMetadata(
+        name="ReasonIR-8B",
+        method="reasonir",
+        description="reasonir/ReasonIR-8B — SOTA reasoning-intensive retriever on the BRIGHT benchmark. No model_id needed.",
+        speed=Speed.VERY_SLOW,
+        accuracy=Accuracy.STATE_OF_THE_ART,
+        gpu_required=True,
+        memory_mb=16000,
+        best_for=["reasoning-intensive queries", "BRIGHT benchmark", "complex multi-hop QA", "science queries"],
+    ),
+    "reason-embed-qwen3-8b": ModelMetadata(
+        name="ReasonEmbed Qwen3-8B",
+        method="reason-embed",
+        description="hanhainebula/reason-embed-qwen3-8b-0928 — Qwen3-8B for reasoning retrieval. Use model_id='qwen3-8b'.",
+        speed=Speed.VERY_SLOW,
+        accuracy=Accuracy.STATE_OF_THE_ART,
+        gpu_required=True,
+        memory_mb=16000,
+        best_for=["reasoning-intensive retrieval", "complex queries"],
+        model_path="qwen3-8b",
+    ),
+    "reason-embed-qwen3-4b": ModelMetadata(
+        name="ReasonEmbed Qwen3-4B",
+        method="reason-embed",
+        description="hanhainebula/reason-embed-qwen3-4b-0928 — balanced Qwen3-4B for reasoning retrieval. Use model_id='qwen3-4b'.",
+        speed=Speed.SLOW,
+        accuracy=Accuracy.EXCELLENT,
+        gpu_required=True,
+        memory_mb=8000,
+        best_for=["reasoning retrieval", "balanced accuracy/speed"],
+        model_path="qwen3-4b",
+    ),
+    "reason-embed-llama-8b": ModelMetadata(
+        name="ReasonEmbed LLaMA-3.1-8B",
+        method="reason-embed",
+        description="hanhainebula/reason-embed-llama-3.1-8b-0928 — LLaMA-3.1-8B for reasoning retrieval. Use model_id='llama-8b'.",
+        speed=Speed.SLOW,
+        accuracy=Accuracy.EXCELLENT,
+        gpu_required=True,
+        memory_mb=16000,
+        best_for=["reasoning retrieval", "open-source LLaMA backbone"],
+        model_path="llama-8b",
+    ),
+    "bge-reasoner-embed": ModelMetadata(
+        name="BGE Reasoner Embed (Qwen3-8B)",
+        method="bge-reasoner-embed",
+        description="BAAI/bge-reasoner-embed-qwen3-8b-0923 — BGE reasoning-augmented retriever. No model_id needed.",
+        speed=Speed.SLOW,
+        accuracy=Accuracy.EXCELLENT,
+        gpu_required=True,
+        memory_mb=16000,
+        best_for=["reasoning-augmented retrieval", "BEIR benchmarks", "complex queries"],
+    ),
 }
 
 
+
+
 # =============================================================================
 # RERANKER REGISTRY
 # =============================================================================