57 lines
2.2 KiB
Python
57 lines
2.2 KiB
Python
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()
|