Spaces:
Running
Running
| """Opt-in on-host three-resident-model CPU/RSS benchmark. | |
| Run from service root with model snapshots downloaded and API_KEY configured: | |
| BENCH_LONG=1 python scripts/bench_local.py | |
| Reports are checkpointed after every case to survive failures/OOM. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import threading | |
| import time | |
| from datetime import datetime, timezone | |
| from pathlib import Path | |
| from app.embedding import MODEL_ID as E5_ID | |
| from app.models import ModelRegistry, QWEN_ID, COLBERT_ID, memory_mb | |
| REPORT = Path(os.getenv("BENCH_REPORT", "bench-result.json")) | |
| LONG = os.getenv("BENCH_LONG", "0") == "1" | |
| rows = [] | |
| report = {"time": datetime.now(timezone.utc).isoformat(), | |
| "long_context": LONG, "startup": None, "cases": rows} | |
| def save(): | |
| REPORT.write_text(json.dumps(report, indent=2, ensure_ascii=False), | |
| encoding="utf-8") | |
| def repeated_text(tokenizer, max_tokens: int, prefix: str = "") -> str: | |
| """Find a repetition count whose true tokenizer length is <= target. | |
| Model-specific input overhead (query instructions/markers) is counted. | |
| A result slightly shorter than target is reported, not mislabeled. | |
| """ | |
| base = "Planet earth history. " | |
| def length(n): | |
| return len(tokenizer(prefix + base * n, add_special_tokens=True, | |
| truncation=False)["input_ids"]) | |
| lo, hi = 0, max_tokens | |
| while lo < hi: | |
| mid = (lo + hi + 1) // 2 | |
| if length(mid) <= max_tokens: | |
| lo = mid | |
| else: | |
| hi = mid - 1 | |
| if lo == 0: | |
| raise RuntimeError("Model prefix exceeds the test target") | |
| return base * lo | |
| def measure(name, runner, token_count=None): | |
| readings = [] | |
| finish = threading.Event() | |
| def monitor(): | |
| while not finish.wait(0.05): | |
| value = memory_mb()["rss_mb"] | |
| if value is not None: | |
| readings.append(value) | |
| sampler = threading.Thread(target=monitor, daemon=True) | |
| sampler.start() | |
| started = time.perf_counter() | |
| result = {"case": name, "requested_tokens": token_count} | |
| try: | |
| payload = runner() | |
| result["status"] = "pass" | |
| result["output_tokens"] = payload.tokens_count | |
| if hasattr(payload, "embedding"): | |
| result["dimensions"] = len(payload.embedding) | |
| else: | |
| result["dimensions"] = len(payload.token_embeddings[0]) | |
| result["vector_count"] = len(payload.token_embeddings) | |
| except Exception as error: | |
| result["status"] = "fail" | |
| result["error"] = f"{type(error).__name__}: {error}" | |
| finally: | |
| finish.set() | |
| sampler.join() | |
| result["elapsed_seconds"] = round(time.perf_counter() - started, 3) | |
| result["rss_mb"] = memory_mb()["rss_mb"] | |
| result["peak_sampled_rss_mb"] = max(readings, default=None) | |
| result["peak_process_rss_mb"] = memory_mb()["peak_rss_mb"] | |
| rows.append(result) | |
| save() | |
| print(result, flush=True) | |
| def main(): | |
| started = time.perf_counter() | |
| registry = ModelRegistry.load() # all 3 weights resident, once | |
| report["startup"] = { | |
| "seconds": round(time.perf_counter() - started, 3), | |
| "per_model_rss": registry.startup_memory, | |
| "post_load_ram": memory_mb(), | |
| } | |
| save() | |
| for kind in ("query", "passage"): | |
| for language, text in ( | |
| ("en", "A long voyage through the forest."), | |
| ("ru", "Долгое путешествие сквозь лес."), | |
| ("he", "מסע ארוך דרך היער."), | |
| ): | |
| measure(f"e5-{kind}-{language}", | |
| lambda t=text, k=kind: registry.e5.encode(t, k)) | |
| measure(f"qwen-{kind}-{language}", | |
| lambda t=text, k=kind: registry.qwen.encode(t, k)) | |
| measure(f"colbert-{kind}-{language}", | |
| lambda t=text, k=kind: registry.colbert.encode(t, k)) | |
| if LONG: | |
| for target in (8192, 32768): | |
| text = repeated_text(registry.qwen.tokenizer, target) | |
| actual = len(registry.qwen.tokenizer( | |
| text, add_special_tokens=True, truncation=False)["input_ids"]) | |
| measure(f"qwen-passage-{target}", lambda t=text: registry.qwen.encode(t, "passage"), actual) | |
| model = registry.colbert.model | |
| prefix = model.document_prefix | |
| text = repeated_text(model.tokenizer, 8192, prefix) | |
| count = len(model.tokenizer( | |
| prefix + text, add_special_tokens=True, truncation=False)["input_ids"]) | |
| measure("colbert-passage-8192", | |
| lambda: registry.colbert.encode(text, "passage"), count) | |
| if __name__ == "__main__": | |
| main() | |