benchmarks: v3.3.0 reproduction results + Ollama rerank backend

benchmarks/locomo_bench.py

@@ -510,19 +510,27 @@ def palace_assign_rooms(sessions, sample_id, api_key, cache, model="claude-haiku
 
 
 def llm_rerank_locomo(
-    question, retrieved_ids, retrieved_docs, api_key, top_k=10, model="claude-sonnet-4-6"
+    question,
+    retrieved_ids,
+    retrieved_docs,
+    api_key,
+    top_k=10,
+    model="claude-sonnet-4-6",
+    backend="anthropic",
+    base_url="",
 ):
     """
     Ask LLM to pick the single most relevant document for this question.
     Returns reordered retrieved_ids with the best candidate first.
+
+    Supports backend="anthropic" (default) or "ollama" (OpenAI-compat endpoint).
     """
     candidates = retrieved_ids[:top_k]
     candidate_docs = retrieved_docs[:top_k]
 
     if len(candidates) <= 1:
         return retrieved_ids
 
-    # Build numbered list of candidates
     lines = []
     for i, (cid, doc) in enumerate(zip(candidates, candidate_docs), 1):
         snippet = doc[:300].replace("\n", " ")
@@ -534,35 +542,51 @@ def llm_rerank_locomo(
         f"Reply with just the number (1-{len(candidates)}).\n\n" + "\n".join(lines)
     )
 
-    payload = json.dumps(
-        {
-            "model": model,
-            "max_tokens": 8,
-            "messages": [{"role": "user", "content": prompt}],
-        }
-    ).encode("utf-8")
-
-    req = urllib.request.Request(
-        "https://api.anthropic.com/v1/messages",
-        data=payload,
-        headers={
+    if backend == "ollama":
+        url = (base_url or "http://localhost:11434").rstrip("/") + "/v1/chat/completions"
+        payload = json.dumps(
+            {
+                "model": model,
+                "messages": [{"role": "user", "content": prompt}],
+                "max_tokens": 1024,
+                "temperature": 0.0,
+            }
+        ).encode("utf-8")
+        headers = {"content-type": "application/json"}
+        if api_key:
+            headers["authorization"] = f"Bearer {api_key}"
+    else:
+        url = "https://api.anthropic.com/v1/messages"
+        payload = json.dumps(
+            {
+                "model": model,
+                "max_tokens": 8,
+                "messages": [{"role": "user", "content": prompt}],
+            }
+        ).encode("utf-8")
+        headers = {
             "x-api-key": api_key,
             "anthropic-version": "2023-06-01",
             "content-type": "application/json",
-        },
-        method="POST",
-    )
+        }
+
+    req = urllib.request.Request(url, data=payload, headers=headers, method="POST")
 
     import socket as _socket
 
     for _attempt in range(3):
         try:
-            with urllib.request.urlopen(req, timeout=30) as resp:
+            with urllib.request.urlopen(req, timeout=120 if backend == "ollama" else 30) as resp:
                 result = json.loads(resp.read())
-            raw = result["content"][0]["text"].strip()
-            m = re.search(r"\b(\d+)\b", raw)
+            if backend == "ollama":
+                msg = result["choices"][0]["message"]
+                raw = (msg.get("content") or "").strip() or (msg.get("reasoning") or "").strip()
+            else:
+                raw = result["content"][0]["text"].strip()
+            # Take LAST integer — reasoning models often count candidates first
+            m = re.search(r"\b(\d+)\b", raw[::-1])
             if m:
-                pick = int(m.group(1))
+                pick = int(m.group(1)[::-1])
                 if 1 <= pick <= len(candidates):
                     chosen_id = candidates[pick - 1]
                     reordered = [chosen_id] + [cid for cid in retrieved_ids if cid != chosen_id]
@@ -608,6 +632,8 @@ def run_benchmark(
     palace_cache_file=None,
     palace_model="claude-haiku-4-5-20251001",
     embed_model="default",
+    llm_backend="anthropic",
+    llm_base_url="",
 ):
     """Run LoCoMo retrieval benchmark."""
     with open(data_file) as f:
@@ -619,8 +645,12 @@ def run_benchmark(
     api_key = ""
     if llm_rerank_enabled or mode == "palace":
         api_key = _load_api_key(llm_key)
-        if not api_key:
-            print(f"ERROR: --mode {mode} requires an API key (--llm-key or ANTHROPIC_API_KEY).")
+        # Ollama backend doesn't require an Anthropic key. Palace mode still does
+        # (it uses Anthropic for room-assignment indexing) — so only relax the
+        # requirement when rerank is the ONLY llm use and backend is ollama.
+        needs_key = mode == "palace" or (llm_rerank_enabled and llm_backend == "anthropic")
+        if needs_key and not api_key:
+            print(f"ERROR: --mode {mode} / --llm-rerank (anthropic) requires an API key.")
             sys.exit(1)
 
     # Palace mode: load or create room assignment cache
@@ -888,6 +918,8 @@ def run_benchmark(
                         api_key,
                         top_k=rerank_pool,
                         model=llm_model,
+                        backend=llm_backend,
+                        base_url=llm_base_url,
                     )
 
                 # Compute recall
@@ -1013,6 +1045,18 @@ def evidence_to_session_ids(evidence):
         help="Model for LLM rerank (default: claude-sonnet-4-6)",
     )
     parser.add_argument("--llm-key", default="", help="API key (or set ANTHROPIC_API_KEY env var)")
+    parser.add_argument(
+        "--llm-backend",
+        choices=["anthropic", "ollama"],
+        default="anthropic",
+        help="Which API for --llm-rerank. 'anthropic' (default) or 'ollama' "
+        "(OpenAI-compat /v1/chat/completions — works for local + Ollama Cloud).",
+    )
+    parser.add_argument(
+        "--llm-base-url",
+        default="",
+        help="Override base URL for --llm-backend ollama. Default: http://localhost:11434.",
+    )
     parser.add_argument(
         "--hybrid-weight",
         type=float,
@@ -1049,4 +1093,6 @@ def evidence_to_session_ids(evidence):
         palace_cache_file=args.palace_cache,
         palace_model=args.palace_model,
         embed_model=args.embed_model,
+        llm_backend=args.llm_backend,
+        llm_base_url=args.llm_base_url,
     )

benchmarks/longmemeval_bench.py

@@ -2763,7 +2763,15 @@ def apply_temporal(fused_dist, timestamp):
 
 
 def llm_rerank(
-    question, rankings, corpus, corpus_ids, api_key, top_k=10, model="claude-haiku-4-5-20251001"
+    question,
+    rankings,
+    corpus,
+    corpus_ids,
+    api_key,
+    top_k=10,
+    model="claude-haiku-4-5-20251001",
+    backend="anthropic",
+    base_url="",
 ):
     """
     Use an LLM to re-rank the top-k retrieved sessions.
@@ -2772,19 +2780,22 @@ def llm_rerank(
     which single session is most relevant to the question. That session
     is promoted to rank 1; the rest stay in their existing order.
 
-    This closes the gap for "preference" and jargon-dense "assistant"
-    failures where the right session is in top-10 semantically but not
-    top-5 — because the semantic gap (battery life ↔ phone hardware) is
-    too large for embeddings to bridge.
+    Supports two backends:
+      - "anthropic": hits https://api.anthropic.com/v1/messages with x-api-key.
+      - "ollama":    hits {base_url}/v1/chat/completions (OpenAI-compat) —
+                     works for local Ollama (default http://localhost:11434)
+                     and Ollama Cloud (:cloud model tags).
 
     Args:
-        question:    The benchmark question string
-        rankings:    Current ranked list of corpus indices (from any mode)
-        corpus:      List of document strings
-        corpus_ids:  List of corpus IDs (parallel to corpus)
-        api_key:     Anthropic API key string
-        top_k:       How many top sessions to send to LLM (default: 10)
-        model:       Claude model ID for reranking (default: haiku)
+        question:   The benchmark question string
+        rankings:   Current ranked list of corpus indices (from any mode)
+        corpus:     List of document strings
+        corpus_ids: List of corpus IDs (parallel to corpus)
+        api_key:    Anthropic API key (only required for backend="anthropic")
+        top_k:      How many top sessions to send to LLM (default: 10)
+        model:      Model id (Claude model for anthropic, e.g. "minimax-m2.7:cloud" for ollama)
+        backend:    "anthropic" or "ollama"
+        base_url:   Override base URL (ollama default: http://localhost:11434)
 
     Returns:
         Reordered rankings list with LLM's best pick promoted to rank 1.
@@ -2796,7 +2807,6 @@ def llm_rerank(
     if not candidates:
         return rankings
 
-    # Format sessions for the prompt — first 500 chars each, labelled 1..N
     session_blocks = []
     for rank, idx in enumerate(candidates):
         text = corpus[idx][:500].replace("\n", " ").strip()
@@ -2813,49 +2823,68 @@ def llm_rerank(
         f"Most relevant session number:"
     )
 
-    payload = json.dumps(
-        {
-            "model": model,
-            "max_tokens": 8,
-            "messages": [{"role": "user", "content": prompt}],
-        }
-    ).encode("utf-8")
-
-    req = urllib.request.Request(
-        "https://api.anthropic.com/v1/messages",
-        data=payload,
-        headers={
+    if backend == "ollama":
+        url = (base_url or "http://localhost:11434").rstrip("/") + "/v1/chat/completions"
+        payload = json.dumps(
+            {
+                "model": model,
+                "messages": [{"role": "user", "content": prompt}],
+                "max_tokens": 1024,
+                "temperature": 0.0,
+            }
+        ).encode("utf-8")
+        headers = {"content-type": "application/json"}
+        if api_key:
+            headers["authorization"] = f"Bearer {api_key}"
+    else:
+        url = "https://api.anthropic.com/v1/messages"
+        payload = json.dumps(
+            {
+                "model": model,
+                "max_tokens": 8,
+                "messages": [{"role": "user", "content": prompt}],
+            }
+        ).encode("utf-8")
+        headers = {
             "x-api-key": api_key,
             "anthropic-version": "2023-06-01",
             "content-type": "application/json",
-        },
-        method="POST",
-    )
+        }
+
+    req = urllib.request.Request(url, data=payload, headers=headers, method="POST")
 
     import socket as _socket
 
     for _attempt in range(3):
         try:
-            with urllib.request.urlopen(req, timeout=20) as resp:
+            with urllib.request.urlopen(req, timeout=120 if backend == "ollama" else 20) as resp:
                 result = json.loads(resp.read())
-            raw = result["content"][0]["text"].strip()
-            # Parse just the first integer from Haiku's response
-            m = re.search(r"\b(\d+)\b", raw)
+            if backend == "ollama":
+                msg = result["choices"][0]["message"]
+                # Reasoning models (e.g. minimax-m2.7) may emit final answer in "content"
+                # or embed it in "reasoning". Try content first, fall back to reasoning.
+                raw = (msg.get("content") or "").strip()
+                if not raw:
+                    raw = (msg.get("reasoning") or "").strip()
+            else:
+                raw = result["content"][0]["text"].strip()
+            m = re.search(
+                r"\b(\d+)\b", raw[::-1]
+            )  # take LAST integer (rerank models often reason first)
             if m:
-                pick = int(m.group(1))
+                pick = int(m.group(1)[::-1])
                 if 1 <= pick <= len(candidates):
                     chosen_idx = candidates[pick - 1]
                     reordered = [chosen_idx] + [i for i in rankings if i != chosen_idx]
                     return reordered
-            break  # Got a response, even if unparseable — don't retry
+            break
         except (_socket.timeout, TimeoutError):
             if _attempt < 2:
                 import time as _time
 
-                _time.sleep(3)  # brief pause then retry
-            # else fall through to return rankings
+                _time.sleep(3)
         except (urllib.error.URLError, KeyError, ValueError, IndexError, OSError):
-            break  # Non-timeout error — fall back immediately
+            break
 
     return rankings
 
@@ -2919,6 +2948,8 @@ def run_benchmark(
     skip_precompute=False,
     split_file=None,
     split_subset=None,
+    llm_backend="anthropic",
+    llm_base_url="",
 ):
     """Run the full benchmark.
 
@@ -2947,10 +2978,14 @@ def run_benchmark(
     api_key = ""
     if llm_rerank_enabled or mode == "diary":
         api_key = _load_api_key(llm_key)
-        if not api_key:
+        # Ollama backend doesn't require an Anthropic API key; a local/cloud Ollama
+        # daemon with the requested model pulled is enough. Diary mode is always anthropic.
+        needs_key = (llm_backend == "anthropic") or (mode == "diary")
+        if needs_key and not api_key:
             print(
-                "ERROR: --llm-rerank / --mode diary requires an API key. "
-                "Set ANTHROPIC_API_KEY or use --llm-key."
+                "ERROR: --llm-rerank (anthropic backend) / --mode diary requires an API key. "
+                "Set ANTHROPIC_API_KEY or use --llm-key. For ollama backend, pass "
+                "--llm-backend ollama."
             )
             sys.exit(1)
 
@@ -3100,7 +3135,15 @@ def run_benchmark(
         if llm_rerank_enabled:
             rerank_pool = 20 if mode in ("hybrid_v3", "hybrid_v4", "palace") else 10
             rankings = llm_rerank(
-                question, rankings, corpus, corpus_ids, api_key, top_k=rerank_pool, model=llm_model
+                question,
+                rankings,
+                corpus,
+                corpus_ids,
+                api_key,
+                top_k=rerank_pool,
+                model=llm_model,
+                backend=llm_backend,
+                base_url=llm_base_url,
             )
 
         # Evaluate at session level
@@ -3276,7 +3319,21 @@ def run_benchmark(
         default="claude-haiku-4-5-20251001",
         help="Model for LLM re-ranking and diary ingest "
         "(default: claude-haiku-4-5-20251001). "
-        "Use 'claude-sonnet-4-6' for Sonnet comparison.",
+        "Use 'claude-sonnet-4-6' for Sonnet comparison. "
+        "With --llm-backend ollama, use an Ollama model tag like 'minimax-m2.7:cloud'.",
+    )
+    parser.add_argument(
+        "--llm-backend",
+        choices=["anthropic", "ollama"],
+        default="anthropic",
+        help="Which API to hit for --llm-rerank. 'anthropic' (default) uses Anthropic's "
+        "/v1/messages endpoint. 'ollama' uses Ollama's OpenAI-compatible "
+        "/v1/chat/completions endpoint (works with local Ollama and Ollama Cloud).",
+    )
+    parser.add_argument(
+        "--llm-base-url",
+        default="",
+        help="Override base URL for --llm-backend ollama. Defaults to http://localhost:11434.",
     )
     parser.add_argument(
         "--diary-cache",
@@ -3380,4 +3437,6 @@ def run_benchmark(
         args.skip_precompute,
         split_file=args.split_file,
         split_subset=split_subset,
+        llm_backend=args.llm_backend,
+        llm_base_url=args.llm_base_url,
     )