101 lines
3.0 KiB
Python
101 lines
3.0 KiB
Python
#!/usr/bin/env python
|
|
"""
|
|
测试文档解析（仅测试3个文件，不调用LLM）
|
|
"""
|
|
|
|
import sys
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
# 修复：移除conda环境的路径，确保使用虚拟环境的包
|
|
sys.path = [p for p in sys.path if 'anaconda3' not in p.lower() and 'conda' not in p.lower()]
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
from src.preprocessing.document_parser import DocumentParser
|
|
from src.utils.config import load_config
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def main():
|
|
"""测试文档解析"""
|
|
logger.info("=" * 60)
|
|
logger.info("测试文档解析（3个文件）")
|
|
logger.info("=" * 60)
|
|
|
|
# 配置
|
|
config = load_config()
|
|
data_dir = Path("../data/1法律")
|
|
output_dir = Path("./output")
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# 解析文档
|
|
logger.info(f"解析目录: {data_dir}")
|
|
parser = DocumentParser()
|
|
docs = parser.parse_directory(str(data_dir))
|
|
|
|
# 只处理前3个
|
|
docs = docs[:3]
|
|
logger.info(f"处理 {len(docs)} 个文档")
|
|
|
|
# 切分为TextUnit
|
|
logger.info("切分文档为TextUnit...")
|
|
all_textunits = []
|
|
for i, doc in enumerate(docs, 1):
|
|
logger.info(f"处理文档 {i}/{len(docs)}: {Path(doc.get('file_path', '')).name}")
|
|
textunits = parser.split_into_textunits(
|
|
doc,
|
|
max_length=config.MAX_TEXTUNIT_LENGTH
|
|
)
|
|
for tu in textunits:
|
|
tu["doc_id"] = doc.get("file_path", "")
|
|
tu["id"] = f"{doc.get('file_path', '')}_{tu.get('paragraph_index', 0)}"
|
|
all_textunits.extend(textunits)
|
|
logger.info(f" - 生成 {len(textunits)} 个TextUnit")
|
|
|
|
logger.info(f"共生成 {len(all_textunits)} 个TextUnit")
|
|
|
|
# 保存结果
|
|
output_file = output_dir / "test_textunits.json"
|
|
result = {
|
|
"num_docs": len(docs),
|
|
"num_textunits": len(all_textunits),
|
|
"docs": [
|
|
{
|
|
"file_path": doc.get("file_path", ""),
|
|
"title": doc.get("title", ""),
|
|
"text_length": len(doc.get("text", "")),
|
|
}
|
|
for doc in docs
|
|
],
|
|
"sample_textunits": [
|
|
{
|
|
"id": tu.get("id", ""),
|
|
"text": tu.get("text", "")[:200] + "..." if len(tu.get("text", "")) > 200 else tu.get("text", ""),
|
|
"paragraph_index": tu.get("paragraph_index", 0),
|
|
}
|
|
for tu in all_textunits[:10] # 只保存前10个作为示例
|
|
]
|
|
}
|
|
|
|
with open(output_file, "w", encoding="utf-8") as f:
|
|
json.dump(result, f, ensure_ascii=False, indent=2)
|
|
|
|
logger.info("=" * 60)
|
|
logger.info("测试完成！")
|
|
logger.info(f"结果已保存到: {output_file}")
|
|
logger.info(f"统计信息:")
|
|
logger.info(f" - 文档数: {result['num_docs']}")
|
|
logger.info(f" - TextUnit数: {result['num_textunits']}")
|
|
logger.info("=" * 60)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|