feat(retrieval): offer device benchmark

shared/retrieval/benchmark.py

@@ -0,0 +1,56 @@
+import argparse
+import json
+import os
+from pathlib import Path
+import statistics
+import sys
+import time
+
+
+def elapsed(call):
+    started = time.monotonic()
+    call()
+    return time.monotonic() - started
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model-cache', required=True, type=Path)
+    parser.add_argument('--max-seconds', type=float, default=float(os.environ.get('QWEN_BENCHMARK_MAX_SECONDS', '8')))
+    args = parser.parse_args()
+    if args.max_seconds <= 0:
+        raise ValueError('--max-seconds must be greater than zero.')
+    os.environ['HF_HOME'] = str(args.model_cache)
+    from server import shared_encoder, shared_reranker
+
+    encoder = shared_encoder()
+    reranker = shared_reranker()
+    query = 'Find the implementation that updates the selected client configuration.'
+    documents = ['Path: scripts/commands/install.py\n\n' + query] * 50
+    print('Benchmark: warming up embedding model...', file=sys.stderr, flush=True)
+    encoder.encode([query], prompt_name='query')
+    print('Benchmark: probing reranker speed...', file=sys.stderr, flush=True)
+    probe = elapsed(lambda: reranker.predict([(query, documents[0])], show_progress_bar=False))
+    if probe * len(documents) > args.max_seconds:
+        print(json.dumps({
+            'embeddingSeconds': None,
+            'rerankingSeconds': round(probe, 3),
+            'maxSeconds': args.max_seconds,
+            'recommended': False,
+        }))
+        return
+    print('Benchmark: measuring full retrieval batch...', file=sys.stderr, flush=True)
+    reranker.predict([(query, document) for document in documents], show_progress_bar=False)
+    embedding = statistics.median(elapsed(lambda: encoder.encode([query], prompt_name='query')) for _ in range(3))
+    reranking = statistics.median(elapsed(lambda: reranker.predict([(query, document) for document in documents], show_progress_bar=False)) for _ in range(3))
+    result = {
+        'embeddingSeconds': round(embedding, 3),
+        'rerankingSeconds': round(reranking, 3),
+        'maxSeconds': args.max_seconds,
+        'recommended': embedding <= args.max_seconds and reranking <= args.max_seconds,
+    }
+    print(json.dumps(result))
+
+
+if __name__ == '__main__':
+    main()