222 lines
9.4 KiB
Python
222 lines
9.4 KiB
Python
"""Reproducible CPU smoke evaluation; no application database or remote inference service."""
|
||
|
||
import argparse
|
||
import hashlib
|
||
import importlib.metadata
|
||
import json
|
||
import math
|
||
import platform
|
||
import statistics
|
||
import sys
|
||
from dataclasses import asdict
|
||
from pathlib import Path
|
||
from time import perf_counter
|
||
|
||
from app.services.embedding import LOCAL_MODELS, EmbeddingError, LocalEmbeddingProvider
|
||
|
||
|
||
def load_dataset(path: Path) -> tuple[dict, str]:
|
||
with path.open("rb") as source:
|
||
raw = source.read(2_000_001)
|
||
if len(raw) > 2_000_000:
|
||
raise ValueError("评估输入不得超过 2MB")
|
||
data = json.loads(raw)
|
||
if not isinstance(data, dict):
|
||
raise ValueError("评估输入须为 JSON 对象")
|
||
documents, queries = data["documents"], data["queries"]
|
||
if not isinstance(documents, list) or not isinstance(queries, list):
|
||
raise ValueError("documents 与 queries 须为列表")
|
||
if not 1 <= len(documents) <= 1000 or not 1 <= len(queries) <= 1000:
|
||
raise ValueError("文档和问题各须为 1–1000 条")
|
||
for rows in [documents, queries]:
|
||
if any(not isinstance(row, dict) for row in rows):
|
||
raise ValueError("每条记录须为对象")
|
||
ids = [row["id"] for row in rows]
|
||
if any(not isinstance(x, str) or not x or len(x) > 100 for x in ids) or len(
|
||
set(ids)
|
||
) != len(ids):
|
||
raise ValueError("每组 ID 必须非空、唯一且最多 100 字符")
|
||
if any(
|
||
not isinstance(row["text"], str) or not row["text"].strip() or len(row["text"]) > 10000
|
||
for row in rows
|
||
):
|
||
raise ValueError("每条文本须为非空且最多 10000 字符")
|
||
ids = {row["id"] for row in documents}
|
||
for query in queries:
|
||
relevant = query["relevant_ids"]
|
||
if (
|
||
not isinstance(relevant, list)
|
||
or not relevant
|
||
or any(not isinstance(x, str) for x in relevant)
|
||
or len(set(relevant)) != len(relevant)
|
||
or not set(relevant) <= ids
|
||
):
|
||
raise ValueError("相关文档标注必须非空、唯一并引用现有文档")
|
||
if not isinstance(data["dataset_id"], str) or not isinstance(data["purpose"], str):
|
||
raise ValueError("必须填写数据集 ID 与用途")
|
||
return data, hashlib.sha256(raw).hexdigest()
|
||
|
||
|
||
def rank_and_measure(documents, queries, document_vectors, query_vectors, *, top_k=3):
|
||
if not documents or not queries or top_k < 1:
|
||
raise ValueError("评估须有文档、问题且 top_k 为正数")
|
||
if len(document_vectors) != len(documents) or len(query_vectors) != len(queries):
|
||
raise ValueError("评估向量数量与输入不一致")
|
||
rows = []
|
||
for query, vector in zip(queries, query_vectors, strict=True):
|
||
ranked = sorted(
|
||
[
|
||
(doc["id"], math.fsum(x * y for x, y in zip(vector, embedding, strict=True)))
|
||
for doc, embedding in zip(documents, document_vectors, strict=True)
|
||
],
|
||
key=lambda pair: (-pair[1], pair[0]),
|
||
)
|
||
relevant = set(query["relevant_ids"])
|
||
first = next(
|
||
(rank for rank, (id_, _) in enumerate(ranked[:top_k], 1) if id_ in relevant), None
|
||
)
|
||
rows.append(
|
||
{
|
||
"query_id": query["id"],
|
||
"top_results": [
|
||
{"id": id_, "score": round(score, 8)} for id_, score in ranked[:top_k]
|
||
],
|
||
"hit_at_1": int(ranked[0][0] in relevant),
|
||
"recall_at_k": len({id_ for id_, _ in ranked[:top_k]} & relevant) / len(relevant),
|
||
"reciprocal_rank_at_k": 1 / first if first else 0,
|
||
}
|
||
)
|
||
return {
|
||
"top_k": top_k,
|
||
"hit_at_1": statistics.mean(row["hit_at_1"] for row in rows),
|
||
"recall_at_k": statistics.mean(row["recall_at_k"] for row in rows),
|
||
"mrr_at_k": statistics.mean(row["reciprocal_rank_at_k"] for row in rows),
|
||
"per_query": rows,
|
||
}
|
||
|
||
|
||
def measure(call, repeats):
|
||
call() # Warm this exact input shape before taking the recorded samples.
|
||
timings = []
|
||
result = None
|
||
for _ in range(repeats):
|
||
start = perf_counter()
|
||
result = call()
|
||
timings.append((perf_counter() - start) * 1000)
|
||
return result, {
|
||
"samples_ms": [round(value, 3) for value in timings],
|
||
"median_ms": round(statistics.median(timings), 3),
|
||
}
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description="本地 CPU Embedding 实测；合成结果不能用于论文结论"
|
||
)
|
||
parser.add_argument("--model", choices=LOCAL_MODELS, required=True)
|
||
parser.add_argument(
|
||
"--dataset", type=Path, default=Path(__file__).with_name("synthetic_zh.json")
|
||
)
|
||
parser.add_argument("--batch-size", type=int, default=8)
|
||
parser.add_argument("--threads", type=int, default=2)
|
||
parser.add_argument("--repeats", type=int, default=5)
|
||
parser.add_argument("--cache-dir", type=Path)
|
||
parser.add_argument(
|
||
"--download", action="store_true", help="允许下载固定 revision 的公开模型文件"
|
||
)
|
||
parser.add_argument("--output", type=Path)
|
||
args = parser.parse_args()
|
||
if sys.platform != "linux":
|
||
parser.error("当前 RSS 测量仅支持 Linux；Provider 可在其他支持的平台单独使用")
|
||
import resource
|
||
|
||
if not 1 <= args.repeats <= 20:
|
||
parser.error("repeats 必须为 1–20")
|
||
dataset, digest = load_dataset(args.dataset)
|
||
print("初始化固定版本的 CPU 模型。", file=sys.stderr, flush=True)
|
||
started = perf_counter()
|
||
provider = LocalEmbeddingProvider(
|
||
LOCAL_MODELS[args.model],
|
||
batch_size=args.batch_size,
|
||
threads=args.threads,
|
||
cache_dir=args.cache_dir,
|
||
offline=not args.download,
|
||
)
|
||
initialization_seconds = perf_counter() - started
|
||
print("模型就绪，执行预热和编码测量。", file=sys.stderr, flush=True)
|
||
provider.encode(["预热测试"], query=True)
|
||
corpus = [row["text"] for row in dataset["documents"]]
|
||
queries = [row["text"] for row in dataset["queries"]]
|
||
documents, corpus_time = measure(lambda: provider.encode(corpus), args.repeats)
|
||
vectors, query_time = measure(lambda: provider.encode(queries, query=True), args.repeats)
|
||
_, single_time = measure(lambda: provider.encode([queries[0]], query=True), args.repeats)
|
||
long_text = "合成测试资料仅用于测量本地编码资源，不表示任何真实景区信息。" * 12
|
||
_, long_time = measure(lambda: provider.encode([long_text] * args.batch_size), args.repeats)
|
||
try:
|
||
provider.encode(["超长测试" * 500])
|
||
except EmbeddingError:
|
||
length_guard = True
|
||
else:
|
||
raise AssertionError("超长输入未被拒绝")
|
||
import torch
|
||
|
||
report = {
|
||
"schema_version": 1,
|
||
"evaluation_kind": "synthetic_smoke"
|
||
if dataset["purpose"] == "synthetic_smoke_not_domain_benchmark"
|
||
else "provided_dataset_not_independently_verified",
|
||
"dataset": {
|
||
"id": dataset["dataset_id"],
|
||
"purpose": dataset["purpose"],
|
||
"sha256": digest,
|
||
"documents": len(corpus),
|
||
"queries": len(queries),
|
||
},
|
||
"embedding_spec": asdict(provider.spec),
|
||
"space_id": provider.spec.space_id,
|
||
"runtime": {
|
||
"python": platform.python_version(),
|
||
"torch": torch.__version__,
|
||
"transformers": importlib.metadata.version("transformers"),
|
||
"huggingface_hub": importlib.metadata.version("huggingface_hub"),
|
||
"numpy": importlib.metadata.version("numpy"),
|
||
"tokenizers": importlib.metadata.version("tokenizers"),
|
||
"safetensors": importlib.metadata.version("safetensors"),
|
||
"device": "cpu",
|
||
"threads": torch.get_num_threads(),
|
||
"interop_threads": torch.get_num_interop_threads(),
|
||
"batch_size": args.batch_size,
|
||
"repeats": args.repeats,
|
||
"warmup_per_shape": 1,
|
||
},
|
||
"initialization_seconds": round(initialization_seconds, 3),
|
||
"timings": {
|
||
"corpus_encode": corpus_time,
|
||
"query_batch_encode": query_time,
|
||
"single_query_encode": single_time,
|
||
"long_text_batch_encode": long_time,
|
||
},
|
||
"long_text_characters": len(long_text),
|
||
"long_text_tokens": len(provider.tokenizer(long_text)["input_ids"]),
|
||
"peak_process_rss_mib": round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024, 1),
|
||
"token_overflow_rejected": length_guard,
|
||
"retrieval": rank_and_measure(dataset["documents"], dataset["queries"], documents, vectors),
|
||
"limitations": [
|
||
"合成小样本仅检查流程，不能代表那拉提领域质量或正式实验分数",
|
||
"计时为本机预热后的编码，包含分词/补齐/归一化/结果转换，不含 API、数据库与网络请求",
|
||
"峰值 RSS 是整个独立进程的高水位，包含解释器和模型；初始化包含缓存检查/下载与首次模型加载",
|
||
],
|
||
}
|
||
encoded = json.dumps(report, ensure_ascii=False, indent=2) + "\n"
|
||
if args.output:
|
||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||
args.output.write_text(encoded, encoding="utf-8")
|
||
args.output.chmod(0o600)
|
||
print("评估完成，报告已写入指定本机位置。", file=sys.stderr)
|
||
else:
|
||
print(encoded, end="")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|