280 lines
10 KiB
Python
280 lines
10 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
|
||
import pytest
|
||
|
||
from app.modules.data_process.algorithms import (
|
||
chunk_unstructured,
|
||
desensitize_pii,
|
||
detect_text_format,
|
||
extract_structured_records,
|
||
generate_standard_records,
|
||
normalize_text,
|
||
parse_text_content,
|
||
record_fingerprint,
|
||
score_quality,
|
||
stable_split,
|
||
)
|
||
|
||
|
||
def test_parse_utf8_json_jsonl_csv_markdown_and_txt() -> None:
|
||
parsed_json = parse_text_content(
|
||
b'\xef\xbb\xbf{"data":[{"name":"\xe5\xbc\xa0\xe4\xb8\x89"}]}',
|
||
filename="records.json",
|
||
)
|
||
assert parsed_json.format == "json"
|
||
assert parsed_json.records == ({"name": "张三"},)
|
||
|
||
parsed_jsonl = parse_text_content('{"id":1}\n\n{"id":2}\n', filename="records.jsonl")
|
||
assert parsed_jsonl.format == "jsonl"
|
||
assert parsed_jsonl.records == ({"id": 1}, {"id": 2})
|
||
|
||
parsed_csv = parse_text_content("name,answer\r\nAlice,yes\r\nBob,no", filename="records.csv")
|
||
assert parsed_csv.format == "csv"
|
||
assert parsed_csv.text == "name,answer\nAlice,yes\nBob,no"
|
||
assert parsed_csv.records[1] == {"name": "Bob", "answer": "no"}
|
||
|
||
parsed_markdown = parse_text_content("# 标题\n\n正文", filename="README.md")
|
||
assert parsed_markdown.format == "markdown"
|
||
assert parsed_markdown.records == ()
|
||
|
||
parsed_txt = parse_text_content("普通文本", filename="note.txt")
|
||
assert parsed_txt.format == "txt"
|
||
assert parsed_txt.text == "普通文本"
|
||
|
||
|
||
def test_invalid_utf8_and_malformed_structured_content_fail_loudly() -> None:
|
||
with pytest.raises(ValueError, match="not valid UTF-8"):
|
||
parse_text_content(b"\xff\xfe", filename="broken.txt")
|
||
with pytest.raises(ValueError, match="invalid JSONL at line 2"):
|
||
extract_structured_records('{"id":1}\nnot-json', "jsonl")
|
||
with pytest.raises(ValueError, match="more fields"):
|
||
extract_structured_records("a,b\n1,2,3", "csv")
|
||
|
||
|
||
def test_detect_format_from_content_and_normalize() -> None:
|
||
assert detect_text_format(text='{"id":1}\n{"id":2}') == "jsonl"
|
||
assert detect_text_format(text="# Heading\ntext") == "markdown"
|
||
assert detect_text_format(text="a,b\n1,2") == "csv"
|
||
assert normalize_text("\ufeffＡＢＣ \r\n第二\x00行\u200b\t \r\n") == "ABC\n第二行"
|
||
|
||
|
||
def test_extract_json_scalar_and_nested_values_are_stable() -> None:
|
||
assert extract_structured_records("[1, true, null]", "json") == [
|
||
{"value": 1},
|
||
{"value": True},
|
||
{"value": None},
|
||
]
|
||
result = extract_structured_records(
|
||
json.dumps({"items": [{"text": " 内容 "}], "ignored": 1}, ensure_ascii=False),
|
||
"json",
|
||
)
|
||
assert result == [{"text": "内容"}]
|
||
|
||
|
||
def test_desensitize_pii_returns_masked_text_and_counts() -> None:
|
||
source = "邮箱 a.user+tag@example.com，手机 +86 13800138000，身份证 11010519491231002X。"
|
||
masked, counts = desensitize_pii(source)
|
||
assert masked == "邮箱 [EMAIL]，手机 [PHONE]，身份证 [ID_CARD]。"
|
||
assert counts == {"email": 1, "phone": 1, "id_card": 1, "total": 3}
|
||
|
||
|
||
@pytest.mark.parametrize("method", ["semantic", "heading", "fixed", "custom"])
|
||
def test_chunk_methods_preserve_offsets_and_always_advance(method: str) -> None:
|
||
text = "# 第一章\n" + "甲。" * 18 + "\n# 第二章\n" + "乙。" * 18
|
||
kwargs = {"custom_delimiter": "\\n"} if method == "custom" else {}
|
||
chunks = chunk_unstructured(
|
||
text,
|
||
method=method, # type: ignore[arg-type]
|
||
chunk_size=12,
|
||
chunk_overlap=2,
|
||
min_chunk_size=4,
|
||
**kwargs,
|
||
)
|
||
assert len(chunks) > 1
|
||
assert all(chunk.content == normalize_text(text)[chunk.start : chunk.end] for chunk in chunks)
|
||
assert all(chunk.end > chunk.start for chunk in chunks)
|
||
assert all(left.start < right.start for left, right in zip(chunks, chunks[1:]))
|
||
assert all(chunk.start_line <= chunk.end_line for chunk in chunks)
|
||
|
||
|
||
def test_fixed_chunk_overlap_is_exact_when_chunks_are_large_enough() -> None:
|
||
text = " ".join(f"token{i}" for i in range(30))
|
||
chunks = chunk_unstructured(
|
||
text,
|
||
method="fixed",
|
||
chunk_size=10,
|
||
chunk_overlap=3,
|
||
min_chunk_size=4,
|
||
)
|
||
first_tokens = chunks[0].content.split()
|
||
second_tokens = chunks[1].content.split()
|
||
assert first_tokens[-3:] == second_tokens[:3]
|
||
assert chunks[0].token_count == 10
|
||
|
||
|
||
def test_chunk_line_numbers_treat_newline_as_previous_line_boundary() -> None:
|
||
chunks = chunk_unstructured(
|
||
"第一行。\n第二行。\n第三行。",
|
||
method="custom",
|
||
chunk_size=8,
|
||
chunk_overlap=0,
|
||
min_chunk_size=2,
|
||
custom_delimiter="\\n",
|
||
)
|
||
assert chunks[0].content.endswith("\n")
|
||
assert chunks[0].start_line == 1
|
||
assert chunks[0].end_line == 1
|
||
assert chunks[1].start_line == 2
|
||
|
||
|
||
def test_heading_and_custom_boundaries_are_respected() -> None:
|
||
heading_text = "前言 " * 8 + "\n# 第二章\n" + "正文 " * 12
|
||
heading_chunks = chunk_unstructured(
|
||
heading_text,
|
||
method="heading",
|
||
chunk_size=20,
|
||
chunk_overlap=0,
|
||
min_chunk_size=4,
|
||
)
|
||
assert "# 第二章" not in heading_chunks[0].content
|
||
assert heading_chunks[1].content.startswith("#")
|
||
|
||
custom_chunks = chunk_unstructured(
|
||
"a b c d <CUT> e f g h i j",
|
||
method="custom",
|
||
chunk_size=8,
|
||
chunk_overlap=0,
|
||
min_chunk_size=2,
|
||
custom_delimiter="<CUT>",
|
||
)
|
||
assert custom_chunks[0].content.endswith("<CUT>")
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("field", "block"),
|
||
[
|
||
(
|
||
"preserve_code_blocks",
|
||
"```python\n" + "\n".join(f"value_{i} = {i}" for i in range(30)) + "\n```",
|
||
),
|
||
(
|
||
"preserve_tables",
|
||
"| 字段 | 说明 |\n| --- | --- |\n"
|
||
+ "\n".join(f"| field_{i} | value_{i} |" for i in range(30)),
|
||
),
|
||
(
|
||
"preserve_lists",
|
||
"\n".join(f"- 第 {i} 项需要完整保留" for i in range(30)),
|
||
),
|
||
],
|
||
)
|
||
def test_markdown_protected_blocks_are_not_split(field: str, block: str) -> None:
|
||
text = "前言。" * 15 + "\n" + block + "\n" + "结尾。" * 40
|
||
chunks = chunk_unstructured(
|
||
text,
|
||
method="fixed",
|
||
chunk_size=40,
|
||
chunk_overlap=0,
|
||
min_chunk_size=10,
|
||
**{field: True},
|
||
)
|
||
assert any(block in chunk.content for chunk in chunks)
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("kwargs", "message"),
|
||
[
|
||
({"chunk_size": 0}, "chunk_size"),
|
||
({"chunk_size": 10, "chunk_overlap": 10}, "chunk_overlap"),
|
||
({"chunk_size": 10, "chunk_overlap": 0, "min_chunk_size": 11}, "min_chunk_size"),
|
||
(
|
||
{"chunk_size": 10, "chunk_overlap": 5, "min_chunk_size": 6},
|
||
"cannot exceed",
|
||
),
|
||
({"method": "custom", "custom_delimiter": ""}, "custom_delimiter"),
|
||
],
|
||
)
|
||
def test_chunk_configuration_validation(kwargs: dict[str, object], message: str) -> None:
|
||
with pytest.raises(ValueError, match=message):
|
||
chunk_unstructured("some text", **kwargs) # type: ignore[arg-type]
|
||
|
||
|
||
def test_quality_scoring_covers_all_dimensions_and_duplicates() -> None:
|
||
valid = {
|
||
"instruction": "如何修改收货地址？",
|
||
"input": "订单尚未发货",
|
||
"output": "可以在订单详情页申请修改收货地址。",
|
||
}
|
||
source = "订单尚未发货时，可以在订单详情页申请修改收货地址。"
|
||
first_score = score_quality(valid, min_output_length=10, source_content=source)
|
||
assert first_score.is_valid
|
||
assert first_score.completeness == 100
|
||
assert first_score.length == 100
|
||
assert first_score.readability >= 90
|
||
assert first_score.relevance >= 70
|
||
assert first_score.duplicate == 100
|
||
|
||
duplicate_score = score_quality(valid, known_fingerprints={first_score.fingerprint})
|
||
assert duplicate_score.duplicate == 0
|
||
assert "duplicate_record" in duplicate_score.flags
|
||
|
||
unrelated_score = score_quality(
|
||
valid,
|
||
min_output_length=10,
|
||
source_content="量子计算使用量子比特处理信息。",
|
||
)
|
||
assert unrelated_score.relevance < first_score.relevance
|
||
assert "low_source_relevance" in unrelated_score.flags
|
||
|
||
invalid_score = score_quality({"instruction": "", "output": "短"}, min_output_length=10)
|
||
assert not invalid_score.is_valid
|
||
assert {"missing_instruction", "output_too_short"}.issubset(invalid_score.flags)
|
||
assert record_fingerprint(valid) == record_fingerprint(dict(reversed(list(valid.items()))))
|
||
|
||
|
||
def test_stable_split_is_reproducible_and_validates_ratios() -> None:
|
||
first = stable_split("record-42", seed="task-1")
|
||
assert stable_split("record-42", seed="task-1") == first
|
||
assert first in {"train", "validation", "test"}
|
||
assert stable_split("record-42", {"train": 100, "validation": 0, "test": 0}) == "train"
|
||
with pytest.raises(ValueError, match="sum to 100"):
|
||
stable_split("record", {"train": 80, "validation": 10, "test": 9})
|
||
|
||
|
||
def test_generate_standard_records_supports_json_qa_and_stable_variants() -> None:
|
||
previews = [
|
||
{
|
||
"id": "preview-json",
|
||
"edited_content": json.dumps(
|
||
{"instruction": "问题", "input": "上下文", "output": "答案"},
|
||
ensure_ascii=False,
|
||
),
|
||
},
|
||
{"id": "preview-qa", "editedContent": "问：如何操作？\n答：按步骤操作。"},
|
||
]
|
||
records = generate_standard_records(
|
||
previews,
|
||
qa_pairs_per_item=2,
|
||
semantic_enrichment=True,
|
||
split={"train": 100, "validation": 0, "test": 0},
|
||
split_seed="task-1",
|
||
)
|
||
assert len(records) == 4
|
||
assert records[0]["instruction"] == "问题"
|
||
assert records[0]["input"] == "上下文"
|
||
assert records[0]["output"] == "答案"
|
||
assert records[1]["instruction"].endswith("问题")
|
||
assert records[2]["instruction"] == "如何操作?"
|
||
assert records[2]["output"] == "按步骤操作。"
|
||
assert all(record["status"] == "valid" for record in records)
|
||
assert all(record["split"] == "train" for record in records)
|
||
assert records == generate_standard_records(
|
||
previews,
|
||
qa_pairs_per_item=2,
|
||
semantic_enrichment=True,
|
||
split={"train": 100, "validation": 0, "test": 0},
|
||
split_seed="task-1",
|
||
)
|