feat: evaluate pinned local embedding models on CPU

This commit is contained in:
2026-10-16 09:11:33 +08:00
parent 1c9b351507
commit cf4e65bb9d
11 changed files with 647 additions and 1 deletions
+1
View File
@@ -0,0 +1 @@
"""Offline, synthetic local-model evaluation; never writes the application database."""
+221
View File
@@ -0,0 +1,221 @@
"""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()
+30
View File
@@ -0,0 +1,30 @@
{
"dataset_id": "synthetic-zh-travel-smoke-v1",
"purpose": "synthetic_smoke_not_domain_benchmark",
"documents": [
{"id":"d01","text":"【合成测试】示例园区的停车场设有接驳车,步行游客可换乘到游客中心。"},
{"id":"d02","text":"【合成测试】示例停车区域按车辆停留时段管理,离场时到出口办理停车手续。"},
{"id":"d03","text":"【合成测试】示例展馆提供轮椅通道和电梯,行动不便的访客可从无障碍入口进入。"},
{"id":"d04","text":"【合成测试】示例演出因降雨取消,已购票观众可沿原购票渠道申请退款。"},
{"id":"d05","text":"【合成测试】示例红色文化展厅设有预约讲解,访客可提前登记并选择导览时段。"},
{"id":"d06","text":"【合成测试】示例资料管理员核查发布机构、来源链接和更新时间,只有审核通过的资料才能成为问答依据。"},
{"id":"d07","text":"【合成测试】示例民宿在前台办理入住,旅客需要出示有效身份证件并确认预订信息。"},
{"id":"d08","text":"【合成测试】示例住宿提供早餐,有素食需求的客人可以提前向接待人员说明。"},
{"id":"d09","text":"【合成测试】示例山间步道遇到恶劣天气暂停通行,游客应查看关闭提示并选择室内活动。"},
{"id":"d10","text":"【合成测试】示例游客服务台统一接收捡到的物品,并登记失物线索协助寻找主人。"},
{"id":"d11","text":"【合成测试】示例停车区设有电动车充电点,司机可按现场设备说明补充电量。"},
{"id":"d12","text":"【合成测试】示例场馆的开放安排以入口公告为准,参观者应提前确认当天是否接待。"}
],
"queries": [
{"id":"q01","text":"没有开车的人怎样从停车区去游客中心?","relevant_ids":["d01"]},
{"id":"q02","text":"同行者坐轮椅,可以从哪里进入展馆?","relevant_ids":["d03"]},
{"id":"q03","text":"演出被雨取消,已经付的钱能退吗?","relevant_ids":["d04"]},
{"id":"q04","text":"想听展厅里的故事讲解,如何安排导览?","relevant_ids":["d05"]},
{"id":"q05","text":"怎么确保问答使用的旅游介绍有可靠出处?","relevant_ids":["d06"]},
{"id":"q06","text":"到了住宿地点办理入住需要准备什么证件?","relevant_ids":["d07"]},
{"id":"q07","text":"早饭不吃肉,应该提前找谁说明?","relevant_ids":["d08"]},
{"id":"q08","text":"天气恶劣时还能在山路上游玩吗?","relevant_ids":["d09"]},
{"id":"q09","text":"在路上捡到游客遗落的东西应交到哪里?","relevant_ids":["d10"]},
{"id":"q10","text":"电动汽车快没电了,到哪儿补充电量?","relevant_ids":["d11"]}
]
}