feat: evaluate pinned local embedding models on CPU
This commit is contained in:
@@ -0,0 +1,161 @@
|
||||
"""Optional local embeddings; importing this module never loads model weights or PyTorch."""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Protocol
|
||||
|
||||
QUERY_INSTRUCTION = "为这个句子生成表示以用于检索相关文章："
|
||||
|
||||
|
||||
class EmbeddingError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class EmbeddingSpec:
|
||||
model: str
|
||||
revision: str
|
||||
dimension: int
|
||||
token_limit: int = 512
|
||||
query_instruction: str = QUERY_INSTRUCTION
|
||||
implementation: str = "hf-cls-l2-cpu-fp32-v1"
|
||||
checkpoint_format: str = "safetensors"
|
||||
|
||||
def __post_init__(self):
|
||||
if not isinstance(self.revision, str) or not re.fullmatch(r"[0-9a-f]{40}", self.revision):
|
||||
raise EmbeddingError("模型必须固定到完整 revision SHA")
|
||||
if (
|
||||
type(self.dimension) is not int
|
||||
or type(self.token_limit) is not int
|
||||
or self.dimension < 1
|
||||
or self.token_limit < 2
|
||||
):
|
||||
raise EmbeddingError("向量维度和 token 上限无效")
|
||||
if not self.model or not isinstance(self.query_instruction, str):
|
||||
raise EmbeddingError("模型名称或查询指令无效")
|
||||
if self.checkpoint_format not in {"safetensors", "pytorch-weights-only"}:
|
||||
raise EmbeddingError("不支持的权重格式")
|
||||
|
||||
@property
|
||||
def space_id(self) -> str:
|
||||
return hashlib.sha256(
|
||||
json.dumps(asdict(self), sort_keys=True, ensure_ascii=False).encode()
|
||||
).hexdigest()
|
||||
|
||||
|
||||
SMALL_ZH = EmbeddingSpec("BAAI/bge-small-zh-v1.5", "7999e1d3359715c523056ef9478215996d62a620", 512)
|
||||
BASE_ZH = EmbeddingSpec(
|
||||
"BAAI/bge-base-zh-v1.5",
|
||||
"f03589ceff5aac7111bd60cfc7d497ca17ecac65",
|
||||
768,
|
||||
checkpoint_format="pytorch-weights-only",
|
||||
)
|
||||
LOCAL_MODELS = {"small": SMALL_ZH, "base": BASE_ZH}
|
||||
|
||||
|
||||
class EmbeddingProvider(Protocol):
|
||||
spec: EmbeddingSpec
|
||||
|
||||
def encode(self, texts: list[str], *, query: bool = False) -> list[list[float]]: ...
|
||||
|
||||
|
||||
def validate_vectors(vectors: list[list[float]], count: int, spec: EmbeddingSpec):
|
||||
if len(vectors) != count:
|
||||
raise EmbeddingError("模型返回数量与输入不一致")
|
||||
for vector in vectors:
|
||||
try:
|
||||
finite = all(type(x) in {float, int} and math.isfinite(x) for x in vector)
|
||||
except (TypeError, OverflowError):
|
||||
finite = False
|
||||
if len(vector) != spec.dimension or not finite:
|
||||
raise EmbeddingError("模型返回维度错误或非有限数值")
|
||||
if not math.isclose(sum(x * x for x in vector), 1, abs_tol=1e-4):
|
||||
raise EmbeddingError("模型输出没有正确 L2 归一化")
|
||||
|
||||
|
||||
class LocalEmbeddingProvider:
|
||||
def __init__(
|
||||
self,
|
||||
spec: EmbeddingSpec,
|
||||
*,
|
||||
batch_size: int = 8,
|
||||
threads: int = 2,
|
||||
cache_dir=None,
|
||||
offline: bool = True,
|
||||
):
|
||||
if (
|
||||
type(batch_size) is not int
|
||||
or type(threads) is not int
|
||||
or not 1 <= batch_size <= 32
|
||||
or not 1 <= threads <= 4
|
||||
):
|
||||
raise EmbeddingError("CPU batch_size 必须为 1–32，threads 必须为 1–4")
|
||||
import torch
|
||||
from huggingface_hub import snapshot_download
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
torch.set_num_threads(threads)
|
||||
self.spec, self.batch_size, self.torch = spec, batch_size, torch
|
||||
path = snapshot_download(
|
||||
repo_id=spec.model,
|
||||
revision=spec.revision,
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=offline,
|
||||
token=False,
|
||||
allow_patterns=[
|
||||
"config.json",
|
||||
"model.safetensors"
|
||||
if spec.checkpoint_format == "safetensors"
|
||||
else "pytorch_model.bin",
|
||||
"tokenizer.json",
|
||||
"tokenizer_config.json",
|
||||
"special_tokens_map.json",
|
||||
"vocab.txt",
|
||||
"added_tokens.json",
|
||||
],
|
||||
)
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||
path, local_files_only=True, trust_remote_code=False
|
||||
)
|
||||
self.model = (
|
||||
AutoModel.from_pretrained(
|
||||
path,
|
||||
local_files_only=True,
|
||||
trust_remote_code=False,
|
||||
use_safetensors=spec.checkpoint_format == "safetensors",
|
||||
weights_only=True,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
.to("cpu")
|
||||
.eval()
|
||||
)
|
||||
if self.model.config.hidden_size != spec.dimension:
|
||||
raise EmbeddingError("模型实际维度与配置不一致")
|
||||
if self.tokenizer.model_max_length != spec.token_limit:
|
||||
raise EmbeddingError("模型 tokenizer 上限与配置不一致")
|
||||
|
||||
def encode(self, texts: list[str], *, query: bool = False) -> list[list[float]]:
|
||||
if any(
|
||||
not isinstance(text, str) or not text.strip() or len(text) > 10_000 for text in texts
|
||||
):
|
||||
raise EmbeddingError("每条输入须为非空文本且最多 10000 字符")
|
||||
torch = self.torch
|
||||
vectors = []
|
||||
for offset in range(0, len(texts), self.batch_size):
|
||||
batch = texts[offset : offset + self.batch_size]
|
||||
if query:
|
||||
batch = [self.spec.query_instruction + text for text in batch]
|
||||
encoded = self.tokenizer(
|
||||
batch, padding=True, truncation=False, return_tensors="pt", verbose=False
|
||||
)
|
||||
if encoded["input_ids"].shape[1] > self.spec.token_limit:
|
||||
raise EmbeddingError("输入超出模型 token 上限；请重新切分，禁止静默截断")
|
||||
with torch.inference_mode():
|
||||
features = self.model(**encoded).last_hidden_state[:, 0]
|
||||
normalized = torch.nn.functional.normalize(features, p=2, dim=1)
|
||||
vectors.extend(normalized.tolist())
|
||||
validate_vectors(vectors, len(texts), self.spec)
|
||||
return vectors
|
||||
@@ -0,0 +1 @@
|
||||
"""Offline, synthetic local-model evaluation; never writes the application database."""
|
||||
@@ -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()
|
||||
@@ -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"]}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
# Optional Linux CPU evaluation environment; not required by FastAPI startup.
|
||||
# Install the CPU torch wheel from its official index first; see the evaluation guide.
|
||||
torch==2.14.1+cpu
|
||||
transformers==5.20.0
|
||||
huggingface-hub==2.2.0
|
||||
numpy==2.4.6
|
||||
tokenizers==0.23.3
|
||||
safetensors==0.8.0
|
||||
@@ -0,0 +1,120 @@
|
||||
"""Lightweight contract/metric tests do not require model packages or network access."""
|
||||
|
||||
import json
|
||||
from dataclasses import FrozenInstanceError, replace
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from app.services.embedding import (
|
||||
SMALL_ZH,
|
||||
EmbeddingError,
|
||||
EmbeddingSpec,
|
||||
LocalEmbeddingProvider,
|
||||
validate_vectors,
|
||||
)
|
||||
from evaluation.local_embeddings import load_dataset, rank_and_measure
|
||||
|
||||
|
||||
def test_model_space_identity_is_immutable_and_includes_semantic_configuration():
|
||||
assert SMALL_ZH.space_id == replace(SMALL_ZH).space_id
|
||||
for change in [
|
||||
{"revision": "a" * 40},
|
||||
{"dimension": 768},
|
||||
{"query_instruction": ""},
|
||||
{"token_limit": 256},
|
||||
{"implementation": "another-pooling"},
|
||||
]:
|
||||
assert replace(SMALL_ZH, **change).space_id != SMALL_ZH.space_id
|
||||
with pytest.raises(FrozenInstanceError):
|
||||
SMALL_ZH.dimension = 768
|
||||
|
||||
|
||||
@pytest.mark.parametrize("revision", ["main", "v1.5", "a" * 39, None])
|
||||
def test_model_requires_full_revision(revision):
|
||||
with pytest.raises(EmbeddingError):
|
||||
EmbeddingSpec("example", revision, 2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"vectors,count",
|
||||
[
|
||||
([[1.0, 0.0]], 2),
|
||||
([[1.0]], 1),
|
||||
([[0.0, 0.0]], 1),
|
||||
([[float("nan"), 0.0]], 1),
|
||||
([[float("inf"), 0.0]], 1),
|
||||
([["invalid", 0.0]], 1),
|
||||
([[True, 0.0]], 1),
|
||||
],
|
||||
)
|
||||
def test_invalid_vectors_cannot_pass_contract(vectors, count):
|
||||
spec = EmbeddingSpec("example", "a" * 40, 2)
|
||||
with pytest.raises(EmbeddingError):
|
||||
validate_vectors(vectors, count, spec)
|
||||
validate_vectors([[1.0, 0.0]], 1, spec)
|
||||
|
||||
|
||||
def test_token_overflow_is_rejected_before_any_model_execution():
|
||||
provider = object.__new__(LocalEmbeddingProvider)
|
||||
provider.spec = SMALL_ZH
|
||||
provider.batch_size = 8
|
||||
provider.torch = None
|
||||
|
||||
class Shape:
|
||||
shape = (1, 513)
|
||||
|
||||
provider.tokenizer = lambda *args, **kwargs: {"input_ids": Shape()}
|
||||
# There is intentionally no model: exceeding the budget must never execute one.
|
||||
with pytest.raises(EmbeddingError, match="token"):
|
||||
provider.encode(["占位文本"])
|
||||
|
||||
|
||||
def test_ranking_metrics_match_hand_calculation_and_ties_are_stable():
|
||||
docs = [{"id": id_} for id_ in ["a", "b", "c"]]
|
||||
queries = [{"id": "q1", "relevant_ids": ["b", "c"]}, {"id": "q2", "relevant_ids": ["c"]}]
|
||||
result = rank_and_measure(docs, queries, [[1, 0], [0, 1], [-1, 0]], [[1, 1], [-1, 0]], top_k=2)
|
||||
assert result["hit_at_1"] == 0.5
|
||||
assert result["recall_at_k"] == 0.75
|
||||
assert result["mrr_at_k"] == 0.75
|
||||
assert [x["id"] for x in result["per_query"][0]["top_results"]] == ["a", "b"]
|
||||
|
||||
|
||||
def test_dataset_validation_rejects_unknown_labels_and_duplicate_ids(tmp_path):
|
||||
data = {
|
||||
"dataset_id": "placeholder",
|
||||
"purpose": "synthetic_smoke_not_domain_benchmark",
|
||||
"documents": [{"id": "d", "text": "占位"}],
|
||||
"queries": [{"id": "q", "text": "测试", "relevant_ids": ["missing"]}],
|
||||
}
|
||||
path = tmp_path / "dataset.json"
|
||||
path.write_text(json.dumps(data))
|
||||
with pytest.raises(ValueError):
|
||||
load_dataset(path)
|
||||
data["queries"][0]["relevant_ids"] = ["d"]
|
||||
data["documents"].append({"id": "d", "text": "重复"})
|
||||
path.write_text(json.dumps(data))
|
||||
with pytest.raises(ValueError):
|
||||
load_dataset(path)
|
||||
|
||||
|
||||
def test_bundled_dataset_is_explicitly_synthetic_and_reproducible():
|
||||
path = Path(__file__).resolve().parents[1] / "evaluation" / "synthetic_zh.json"
|
||||
dataset, digest = load_dataset(path)
|
||||
assert dataset["purpose"] == "synthetic_smoke_not_domain_benchmark"
|
||||
assert all(row["text"].startswith("【合成测试】") for row in dataset["documents"])
|
||||
assert load_dataset(path)[1] == digest
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"options", [{"batch_size": 0}, {"batch_size": 33}, {"threads": 5}, {"batch_size": 1.5}]
|
||||
)
|
||||
def test_invalid_resource_options_fail_without_loading_optional_dependencies(options):
|
||||
with pytest.raises(EmbeddingError):
|
||||
LocalEmbeddingProvider(SMALL_ZH, **options)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint", ["pickle", "automatic-fallback"])
|
||||
def test_unsupported_checkpoint_format_is_rejected(checkpoint):
|
||||
with pytest.raises(EmbeddingError):
|
||||
replace(SMALL_ZH, checkpoint_format=checkpoint)
|
||||
Reference in New Issue
Block a user