iomgaa
6bdb802f01
- Add all claude skills (brainstorming, commit, debugging, TDD, etc.) - Add claude hooks (pre-commit-guard, post-edit-quality) - Add research templates (experiment plan, research brief, etc.) - Add claude tools (arxiv/semantic_scholar/openalex fetch, wiki, exa) - Add TRM4 reference implementation as algorithm fidelity baseline - Add research-wiki content (plans, index, graph, query_pack) - Update .gitignore to exclude .graphify_version runtime state
200 lines
7.2 KiB
Python
200 lines
7.2 KiB
Python
"""
|
||
嵌入服务模块单元测试
|
||
====================
|
||
覆盖: 本地模式（embed/embed_tensor/归一化/dim/冻结）、远程模式（真实 API 调用）。
|
||
|
||
本地测试使用轻量模型 all-MiniLM-L6-v2 (dim=384) 加速。
|
||
远程测试使用 .env 中配置的真实 API（需有效密钥）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import numpy as np
|
||
import pytest
|
||
import torch
|
||
|
||
from video_tree_trm.config import Config, EmbedConfig
|
||
from video_tree_trm.embeddings import EmbeddingModel
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Fixtures
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_LOCAL_CONFIG = EmbedConfig(
|
||
backend="local",
|
||
model_name="all-MiniLM-L6-v2",
|
||
embed_dim=384,
|
||
device="cpu",
|
||
api_key="",
|
||
api_url="",
|
||
)
|
||
|
||
|
||
@pytest.fixture(scope="module")
|
||
def local_model() -> EmbeddingModel:
|
||
"""本地嵌入模型（模块级缓存，避免重复加载）。"""
|
||
return EmbeddingModel(_LOCAL_CONFIG)
|
||
|
||
|
||
@pytest.fixture(scope="module")
|
||
def remote_model(real_config: Config) -> EmbeddingModel:
|
||
"""远程嵌入模型（模块级缓存），使用真实 API。"""
|
||
return EmbeddingModel(real_config.embed)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 测试: 本地模式
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestLocalEmbed:
|
||
"""本地 sentence-transformers 后端测试。"""
|
||
|
||
def test_embed_single_text(self, local_model: EmbeddingModel) -> None:
|
||
"""单条文本，验证形状 [1, D]。"""
|
||
result = local_model.embed("hello world")
|
||
assert isinstance(result, np.ndarray)
|
||
assert result.shape == (1, 384)
|
||
|
||
def test_embed_batch(self, local_model: EmbeddingModel) -> None:
|
||
"""批量文本，验证形状 [N, D]。"""
|
||
texts = ["hello", "world", "test"]
|
||
result = local_model.embed(texts)
|
||
assert result.shape == (3, 384)
|
||
|
||
def test_embed_tensor_type(self, local_model: EmbeddingModel) -> None:
|
||
"""验证 embed_tensor() 返回 Tensor。"""
|
||
result = local_model.embed_tensor("hello")
|
||
assert isinstance(result, torch.Tensor)
|
||
assert result.shape == (1, 384)
|
||
assert result.dtype == torch.float32
|
||
|
||
def test_embeddings_normalized(self, local_model: EmbeddingModel) -> None:
|
||
"""验证 L2 范数 ≈ 1.0。"""
|
||
result = local_model.embed(["hello", "world"])
|
||
norms = np.linalg.norm(result, axis=1)
|
||
np.testing.assert_allclose(norms, 1.0, atol=1e-5)
|
||
|
||
def test_dim_property(self, local_model: EmbeddingModel) -> None:
|
||
"""验证 dim 属性一致。"""
|
||
assert local_model.dim == 384
|
||
|
||
def test_local_model_frozen(self, local_model: EmbeddingModel) -> None:
|
||
"""本地模式参数 requires_grad=False。"""
|
||
for param in local_model._model.parameters():
|
||
assert not param.requires_grad
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 测试: 远程模式（真实 API）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestRemoteEmbed:
|
||
"""远程 OpenAI 兼容 API 后端测试（真实 API 调用）。"""
|
||
|
||
def test_remote_embed_single(self, remote_model: EmbeddingModel) -> None:
|
||
"""单条中文文本，验证形状 [1, D]。"""
|
||
result = remote_model.embed("你好世界")
|
||
assert isinstance(result, np.ndarray)
|
||
assert result.shape == (1, remote_model.dim)
|
||
|
||
def test_remote_embed_batch(self, remote_model: EmbeddingModel) -> None:
|
||
"""5 条文本，验证形状 [5, D]。"""
|
||
texts = ["机器学习", "深度学习", "自然语言处理", "计算机视觉", "强化学习"]
|
||
result = remote_model.embed(texts)
|
||
assert result.shape == (5, remote_model.dim)
|
||
|
||
def test_remote_embed_normalized(self, remote_model: EmbeddingModel) -> None:
|
||
"""验证 L2 范数 ≈ 1.0。"""
|
||
result = remote_model.embed(["归一化测试文本", "另一段测试文本"])
|
||
norms = np.linalg.norm(result, axis=1)
|
||
np.testing.assert_allclose(norms, 1.0, atol=1e-5)
|
||
|
||
def test_remote_embed_tensor(self, remote_model: EmbeddingModel) -> None:
|
||
"""验证返回 torch.Tensor + float32。"""
|
||
result = remote_model.embed_tensor("张量测试")
|
||
assert isinstance(result, torch.Tensor)
|
||
assert result.dtype == torch.float32
|
||
assert result.shape == (1, remote_model.dim)
|
||
|
||
def test_remote_dim(self, remote_model: EmbeddingModel, real_config: Config) -> None:
|
||
"""验证 dim 属性与配置一致。"""
|
||
assert remote_model.dim == real_config.embed.embed_dim
|
||
|
||
def test_remote_semantic_similarity(self, remote_model: EmbeddingModel) -> None:
|
||
"""语义相近文本相似度 > 语义无关文本。"""
|
||
# 语义相近对
|
||
vec_cat = remote_model.embed("猫咪在沙发上睡觉")
|
||
vec_kitten = remote_model.embed("小猫趴在沙发上打盹")
|
||
# 语义无关
|
||
vec_math = remote_model.embed("二次方程的求解公式")
|
||
|
||
sim_close = float(np.dot(vec_cat[0], vec_kitten[0]))
|
||
sim_far = float(np.dot(vec_cat[0], vec_math[0]))
|
||
assert sim_close > sim_far, (
|
||
f"语义相近对相似度 ({sim_close:.4f}) 应大于语义无关对 ({sim_far:.4f})"
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 测试: 配置校验
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestConfigValidation:
|
||
"""配置校验测试。"""
|
||
|
||
def test_invalid_backend(self) -> None:
|
||
"""无效 backend 应抛出 ValueError。"""
|
||
config = EmbedConfig(
|
||
backend="invalid",
|
||
model_name="test",
|
||
embed_dim=384,
|
||
device="cpu",
|
||
api_key="",
|
||
api_url="",
|
||
)
|
||
with pytest.raises(ValueError, match="backend"):
|
||
EmbeddingModel(config)
|
||
|
||
def test_remote_missing_api_key(self) -> None:
|
||
"""远程模式缺少 api_key 应抛出 ValueError。"""
|
||
config = EmbedConfig(
|
||
backend="remote",
|
||
model_name="test",
|
||
embed_dim=384,
|
||
device="cpu",
|
||
api_key="",
|
||
api_url="http://localhost:8080/v1",
|
||
)
|
||
with pytest.raises(ValueError, match="api_key"):
|
||
EmbeddingModel(config)
|
||
|
||
def test_remote_missing_api_url(self) -> None:
|
||
"""远程模式缺少 api_url 应抛出 ValueError。"""
|
||
config = EmbedConfig(
|
||
backend="remote",
|
||
model_name="test",
|
||
embed_dim=384,
|
||
device="cpu",
|
||
api_key="test-key",
|
||
api_url="",
|
||
)
|
||
with pytest.raises(ValueError, match="api_url"):
|
||
EmbeddingModel(config)
|
||
|
||
def test_local_dim_mismatch(self) -> None:
|
||
"""本地模式维度不一致应抛出 ValueError。"""
|
||
config = EmbedConfig(
|
||
backend="local",
|
||
model_name="all-MiniLM-L6-v2",
|
||
embed_dim=999, # 实际是 384
|
||
device="cpu",
|
||
api_key="",
|
||
api_url="",
|
||
)
|
||
with pytest.raises(ValueError, match="维度"):
|
||
EmbeddingModel(config)
|