96 lines
6.6 KiB
Python
96 lines
6.6 KiB
Python
"""Frozen tiny code-retrieval regression; genuine Google worker, CPU only.
|
|||
|
|
|
||
|
|
Generated public fixtures, not a production accuracy benchmark. Run from repo
|
||
|
|
after rebuilding CLI and CPU worker. No private folders are indexed.
|
||
|
|
"""
|
||
|
|
import hashlib
|
||
|
|
from datetime import datetime, timezone
|
||
|
|
import importlib.util
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
import subprocess
|
||
|
|
|
||
|
|
ROOT = Path(__file__).resolve().parents[2]
|
||
|
|
SPEC = importlib.util.spec_from_file_location("benchmark", ROOT / "tools/embeddinggemma2/benchmark.py")
|
||
|
|
BENCH = importlib.util.module_from_spec(SPEC)
|
||
|
|
SPEC.loader.exec_module(BENCH)
|
||
|
|
|
||
|
|
FILES = {
|
||
|
|
"a.py": "import hashlib\ndef process(path):\n h = hashlib.sha256()\n with open(path, 'rb') as stream:\n for chunk in iter(lambda: stream.read(65536), b''):\n h.update(chunk)\n return h.hexdigest()\n",
|
||
|
|
"b.cpp": "#include <vector>\nint locate(const std::vector<int>& v, int x) {\n int lo = 0, hi = v.size();\n while (lo < hi) {\n int mid = lo + (hi-lo)/2;\n if (v[mid] < x) lo = mid+1; else hi = mid;\n }\n return lo;\n}\n",
|
||
|
|
"c.ts": "export function process(rows: {category: string; amount: number}[]) {\n const totals = new Map<string, number>();\n for (const row of rows) totals.set(row.category, (totals.get(row.category) ?? 0) + row.amount);\n return totals;\n}\n",
|
||
|
|
"d.py": "import time\nimport urllib.request\ndef process(url):\n for attempt in range(5):\n try:\n return urllib.request.urlopen(url, timeout=10).read()\n except OSError:\n if attempt == 4: raise\n time.sleep(0.2 * 2 ** attempt)\n",
|
||
|
|
"e.cs": "using System;\nusing System.Threading;\nusing System.Threading.Tasks;\npublic static class Module {\n public static async Task Process(SemaphoreSlim gate, Func<Task> work) {\n await gate.WaitAsync();\n try { await work(); } finally { gate.Release(); }\n }\n}\n",
|
||
|
|
"f.php": "<?php\nfunction process($value) {\n return htmlspecialchars($value, ENT_QUOTES | ENT_SUBSTITUTE, 'UTF-8');\n}\n",
|
||
|
|
"g.py": "def process(rows):\n return sorted(rows, key=lambda row: row['price'], reverse=True)[:10]\n",
|
||
|
|
"h.rs": "fn refreshAccessToken(previous: &str) -> String {\n format!(\"renewed:{}\", previous)\n}\n",
|
||
|
|
"notes.txt": "refreshAccessToken retry failed requests with exponential backoff compute a file digest sum amounts by category escape HTML sorted numbers",
|
||
|
|
}
|
||
|
|
QUERIES = [
|
||
|
|
("calculate a file's SHA256 without loading the entire file into memory", "a.py"),
|
||
|
|
("find the first position where sorted numbers are greater than or equal to a target", "b.cpp"),
|
||
|
|
("sum transaction amounts separately for each category", "c.ts"),
|
||
|
|
("retry failed requests with exponential backoff", "d.py"),
|
||
|
|
("release a semaphore even if the asynchronous operation throws", "e.cs"),
|
||
|
|
("escape unsafe HTML characters before displaying user input", "f.php"),
|
||
|
|
("return the ten most expensive records", "g.py"),
|
||
|
|
("Datei prüfen ohne sie vollständig in den Speicher zu laden", "a.py"),
|
||
|
|
("refreshAccessToken", "h.rs"),
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
output = ROOT / ".searchi/code-eval"
|
||
|
|
corpus = output / "corpus"
|
||
|
|
corpus.mkdir(parents=True, exist_ok=True)
|
||
|
|
for name, content in FILES.items():
|
||
|
|
(corpus / name).write_text(content, encoding="utf-8", newline="\n")
|
||
|
|
cli = ROOT / "target/release" / ("searchi.exe" if os.name == "nt" else "searchi")
|
||
|
|
models = ROOT / ".searchi/models/embeddinggemma2"
|
||
|
|
worker = ROOT / ".searchi/runtime/cpu/Release" / ("searchi-gemma-worker.exe" if os.name == "nt" else "searchi-gemma-worker")
|
||
|
|
db = output / "index.db"
|
||
|
|
environment = dict(os.environ, SEARCHI_FORCE_CPU="1")
|
||
|
|
def run(*args):
|
||
|
|
subprocess.run([str(cli), *map(str, args)], env=environment, check=True)
|
||
|
|
run("index", "--db", db, "--root", corpus)
|
||
|
|
run("embed", "--db", db, "--model-dir", models, "--embedding-worker", worker, "--text-only")
|
||
|
|
rows = []
|
||
|
|
for mode in ("lexical_only", "hybrid"):
|
||
|
|
command = [str(cli), "mcp", "--db", str(db)]
|
||
|
|
if mode == "hybrid":
|
||
|
|
command += ["--model-dir", str(models), "--embedding-worker", str(worker)]
|
||
|
|
session = BENCH.Session(command)
|
||
|
|
try:
|
||
|
|
session.mcp(1, "initialize", {"protocolVersion": "2025-03-26", "capabilities": {}, "clientInfo": {"name": "code-eval", "version": "1"}})
|
||
|
|
session.send({"jsonrpc": "2.0", "method": "notifications/initialized"})
|
||
|
|
for number, (query, expected) in enumerate(QUERIES, 2):
|
||
|
|
result, duration = session.mcp(number, "tools/call", {"name": "search", "arguments": {"query": query, "code_only": True, "limit": 10}})
|
||
|
|
assert not result.get("isError"), result
|
||
|
|
data = result["structuredContent"]
|
||
|
|
assert data["mode"] == mode and data["code_only"]
|
||
|
|
names = [r["filename"] for r in data["results"]]
|
||
|
|
assert "notes.txt" not in names
|
||
|
|
rank = names.index(expected) + 1 if expected in names else None
|
||
|
|
source_verified = False
|
||
|
|
if rank is not None:
|
||
|
|
matched = data["results"][rank - 1]
|
||
|
|
source, _ = session.mcp(number + 100, "tools/call", {"name": "read_source", "arguments": {"handle": matched["source_handle"]}})
|
||
|
|
assert not source.get("isError"), source
|
||
|
|
source = source["structuredContent"]
|
||
|
|
assert source["source"] == {"kind": "line", "number": 1}
|
||
|
|
assert source["text"] == FILES[expected]
|
||
|
|
source_verified = True
|
||
|
|
rows.append({"mode": mode, "query": query, "expected": expected, "rank": rank, "results": names, "roundtrip_ms": duration, "source_verified": source_verified})
|
||
|
|
finally:
|
||
|
|
session.close()
|
||
|
|
report = {"recorded_at": datetime.now(timezone.utc).isoformat(), "device": "cpu", "cpu_threads": 1, "model": "google/embeddinggemma-2", "precision": "verified Q8_0 weights, CPU inference", "query_task": "task: code retrieval | query: ", "scope": "8 tiny source files and 1 ordinary-document hard negative; synthetic regression, not general code accuracy", "fixture_sha256": {n: hashlib.sha256(c.encode()).hexdigest() for n, c in FILES.items()}, "queries": rows}
|
||
|
|
(output / "results.json").write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
|
||
|
|
for mode in ("lexical_only", "hybrid"):
|
||
|
|
selected = [r for r in rows if r["mode"] == mode]
|
||
|
|
print(mode, "top1", sum(r["rank"] == 1 for r in selected), "/", len(selected), "top3", sum(r["rank"] is not None and r["rank"] <= 3 for r in selected))
|
||
|
|
print("Recorded", output / "results.json")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|