test(semantic_chunking): replace mock embedding with real model
Use a real embedding model from SiliconFlow for more realistic testing. Update test assertions to validate LaTeX formula preservation and semantic clustering of distinct chapters.
This commit is contained in:
parent
6934e19030
commit
a2ed3d9a10
@ -1,71 +1,77 @@
|
||||
import pytest
|
||||
import numpy as np
|
||||
import os
|
||||
|
||||
# 现在可以安全地导入了,因为顶层不再有重型依赖
|
||||
from yuxi.knowledge.chunking.ragflow_like.parsers import semantic
|
||||
from yuxi.models import select_embedding_model
|
||||
|
||||
@pytest.fixture
|
||||
def mock_embed_fn():
|
||||
"""模拟一个简单的嵌入函数"""
|
||||
def encode(sentences):
|
||||
if isinstance(sentences, str):
|
||||
sentences = [sentences]
|
||||
# 返回模拟的向量 (N, 384)
|
||||
return np.random.rand(len(sentences), 384).astype(np.float32)
|
||||
return encode
|
||||
def embed_fn():
|
||||
"""使用真实的嵌入函数 (从 SiliconFlow 获取)"""
|
||||
model_id = "siliconflow/Qwen/Qwen3-Embedding-0.6B"
|
||||
try:
|
||||
model = select_embedding_model(model_id)
|
||||
def encode(sentences):
|
||||
if isinstance(sentences, str):
|
||||
sentences = [sentences]
|
||||
return model.encode(sentences)
|
||||
return encode
|
||||
except Exception as e:
|
||||
pytest.skip(f"无法初始化真实嵌入模型 {model_id}: {e}")
|
||||
|
||||
@pytest.fixture
|
||||
def sample_markdown():
|
||||
return """
|
||||
# 核心业务说明
|
||||
|
||||
本节主要介绍公司的核心业务流程。
|
||||
|
||||
## 第一部分:研发流程
|
||||
1. 需求分析:收集用户反馈。
|
||||
2. 系统设计:架构方案评审。
|
||||
3. 编码实现:遵循规范。
|
||||
|
||||
## 第二部分:交付标准
|
||||
| 阶段 | 交付物 | 责任人 |
|
||||
| --- | --- | --- |
|
||||
| 交付1 | 源码包 | 开发 |
|
||||
| 交付2 | 测试报告 | 测试 |
|
||||
|
||||
这里是一些额外的语义内容。
|
||||
语义内容应当被聚类在一起。
|
||||
这也是语义内容。
|
||||
"""
|
||||
|
||||
def test_semantic_chunking_basic(mock_embed_fn, sample_markdown):
|
||||
"""测试基本的语义切分逻辑 (使用注入方式)"""
|
||||
with open("test/resource/test4.md", "r", encoding="utf-8") as f:
|
||||
return f.read()
|
||||
def test_semantic_chunking_basic(embed_fn, sample_markdown):
|
||||
"""测试基本的语义切分逻辑 (使用真实嵌入模型)"""
|
||||
|
||||
# 配置切分参数
|
||||
parser_config = {
|
||||
"chunk_token_num": 100, # 触发语义切分
|
||||
"chunk_token_num": 200, # 适中的 token 数以触发更多切分
|
||||
"overlapped_percent": 0
|
||||
}
|
||||
|
||||
# 直接注入 mock_embed_fn,完全跳过 config 加载
|
||||
# 直接注入 embed_fn
|
||||
chunks = semantic.chunk_markdown(
|
||||
sample_markdown,
|
||||
parser_config=parser_config,
|
||||
embed_fn=mock_embed_fn
|
||||
embed_fn=embed_fn
|
||||
)
|
||||
|
||||
# 验证结果
|
||||
assert isinstance(chunks, list)
|
||||
assert len(chunks) > 0
|
||||
|
||||
# 验证标题路径增强是否生效
|
||||
has_enhanced_title = any("# 核心业务说明|第一部分:研发流程" in chunk for chunk in chunks)
|
||||
assert has_enhanced_title, "标题路径增强失效"
|
||||
print(f"\n成功切分为 {len(chunks)} 个片段:")
|
||||
for i, chunk in enumerate(chunks):
|
||||
print(f"--- 片段 {i+1} ---")
|
||||
print(chunk)
|
||||
|
||||
# 验证表格是否被正确保留
|
||||
has_table = any("交付1" in chunk and "源码包" in chunk for chunk in chunks)
|
||||
assert has_table, "表格内容丢失"
|
||||
# 验证标题路径增强是否生效 (检查是否包含 Part 或特殊元素标记)
|
||||
has_enhanced_title = any("|Part" in chunk or "|Math Block" in chunk for chunk in chunks)
|
||||
assert has_enhanced_title, "标题路径增强或分片标记失效"
|
||||
|
||||
print(f"\n成功切分为 {len(chunks)} 个片段")
|
||||
# 验证是否包含特殊的 LaTeX 公式或符号 (匹配输出中的格式)
|
||||
has_formula = any("f _ {0, 1}" in chunk or "\\sigma_ {y, d}" in chunk or "A _ {\\mathrm {s p}}" in chunk for chunk in chunks)
|
||||
assert has_formula, "LaTeX 公式/符号在切分中丢失"
|
||||
|
||||
# 验证语义聚类是否大概生效 (钢制塔架和混凝土塔架内容应该分开)
|
||||
# 我们检查是否有钢制塔架相关的片段和混凝土塔架相关的片段
|
||||
# 使用 startswith 来更准确地定位章节开始,避免匹配到目次/前言中的文字
|
||||
has_steel = any(chunk.startswith("# 6 钢制塔架") for chunk in chunks)
|
||||
has_concrete = any(chunk.startswith("# 7 混凝土塔架") for chunk in chunks)
|
||||
|
||||
assert has_steel and has_concrete, "语义内容丢失 (钢制或混凝土塔架章节开始部分)"
|
||||
|
||||
# 验证它们是否在不同的片段中
|
||||
steel_chunks = [i for i, c in enumerate(chunks) if c.startswith("# 6 钢制塔架")]
|
||||
concrete_chunks = [i for i, c in enumerate(chunks) if c.startswith("# 7 混凝土塔架")]
|
||||
|
||||
# 理论上语义聚类会将这些大章节分开
|
||||
if steel_chunks and concrete_chunks:
|
||||
assert not set(steel_chunks).intersection(set(concrete_chunks)), "钢制塔架和混凝土塔架大章节内容被错误地聚类在同一个 chunk 中了"
|
||||
|
||||
def test_heading_inference():
|
||||
"""测试标题层级推断工具类"""
|
||||
|
||||
Loading…
Reference in New Issue
Block a user