FlashRank Python 教程：候选文档重排、中文评估与延迟测试

python3.11 -m venv .venv
source .venv/bin/activate
python -m pip install 'flashrank==0.2.10'
python -m pip freeze > requirements.lock.txt

"""FlashRank 0.2.10 pairwise/ONNX example; synthetic, public demo data only."""
import copy
import json
import math
import os
import time
from importlib.metadata import version

from flashrank import Ranker, RerankRequest
from tokenizers import Tokenizer


def checked_rerank(ranker, query, candidates, tenant, top_k=2):
    if not isinstance(query, str) or not query.strip():
        raise ValueError('query must be nonempty')
    if not isinstance(tenant, str) or not tenant.strip():
        raise ValueError('tenant must come from authenticated server context')
    if type(top_k) is not int or not 1 <= top_k <= 10:
        raise ValueError('top_k must be between 1 and 10')
    if not 1 <= len(candidates) <= 30:
        raise ValueError('candidate count must be between 1 and 30')
    if ranker.llm_model is not None:
        raise ValueError('this example supports only pairwise ONNX models')
    # Clone the actual model tokenizer and disable its automatic truncation.
    counter = Tokenizer.from_str(ranker.tokenizer.to_str())
    counter.no_truncation()
    counter.no_padding()
    limit = ranker.tokenizer.truncation['max_length']
    originals = {}
    for p in candidates:
        if not isinstance(p, dict):
            raise ValueError('invalid candidate')
        cid, text, meta = p.get('id'), p.get('text'), p.get('meta')
        if not isinstance(cid, str) or not cid or cid in originals:
            raise ValueError('candidate IDs must be nonempty and unique')
        if not isinstance(text, str) or not text.strip():
            raise ValueError('candidate text must be nonempty')
        if not isinstance(meta, dict) or meta.get('tenant_id') != tenant:
            raise ValueError('tenant mismatch; reject before model execution')
        if not meta.get('source') or not meta.get('version'):
            raise ValueError('source and version required')
        if len(counter.encode(query, text).ids) > limit:
            raise ValueError('query-document pair exceeds model token budget')
        originals[cid] = p
    started = time.perf_counter()
    ranked = ranker.rerank(RerankRequest(query=query, passages=copy.deepcopy(candidates)))
    elapsed_ms = (time.perf_counter() - started) * 1000
    if len(ranked) != len(candidates) or {p['id'] for p in ranked} != set(originals):
        raise RuntimeError('unexpected output candidate set')
    for p in ranked:
        original = originals[p['id']]
        if p['text'] != original['text'] or p['meta'] != original['meta']:
            raise RuntimeError('source data changed')
        p['score'] = float(p['score'])
        if not math.isfinite(p['score']):
            raise RuntimeError('non-finite relevance score')
    if any(a['score'] < b['score'] for a, b in zip(ranked, ranked[1:])):
        raise RuntimeError('unexpected score ordering')
    return ranked[:top_k], elapsed_ms


def fixtures():
    texts = [
        ('refund', '退款规则：购买后七天内且未使用服务可申请退款。'),
        ('password', '登录失败时，可以通过注册邮箱重置密码。'),
        ('billing', '发票将在付款成功后发送到客户邮箱。'),
        ('cancel', '关闭自动续费可避免下期扣款，但不会自动退还已付费用。'),
    ]
    return [dict(id=cid, text=text, meta=dict(tenant_id='demo', source='public-demo-policy', version='v1')) for cid, text in texts]


if __name__ == '__main__':
    model = os.environ.get('FLASHRANK_MODEL', 'ms-marco-MultiBERT-L-12')
    start = time.perf_counter()
    ranker = Ranker(model_name=model, cache_dir=os.environ.get('FLASHRANK_CACHE', './model-cache'), max_length=512)
    init_ms = (time.perf_counter() - start) * 1000
    candidates = fixtures()
    before = copy.deepcopy(candidates)
    query = '购买后什么情况下可以退款？'
    rows, first_ms = checked_rerank(ranker, query, candidates, 'demo')
    timings = [checked_rerank(ranker, query, candidates, 'demo')[1] for _ in range(5)]
    assert candidates == before
    print(json.dumps(dict(flashrank=version('flashrank'), model=model, initialization_ms=round(init_ms, 2), first_rerank_ms=round(first_ms, 2), warm_rerank_ms=[round(t, 2) for t in timings], results=rows, input_unchanged=True), ensure_ascii=False, indent=2))

export FLASHRANK_CACHE="$PWD/model-cache"
python rerank_demo.py > result.json