2020年药典学习python调整

data-pipeline/pdf2020_ingest.py

@@ -10,6 +10,7 @@ import httpx
 import fitz  # pymupdf
 from sqlalchemy.ext.asyncio import create_async_engine
 from sqlalchemy import text
+from processors.chunker import DrugChunker
 
 # ============================================
 # 配置
@@ -215,25 +216,43 @@ async def get_embeddings(texts: list[str], text_type: str = "document") -> list[
 
 
 async def ingest_entries(entries: list[dict]):
-    """向量化 + 入库"""
+    """向量化 + 入库。长 section 用 DrugChunker 切分。"""
     engine = create_async_engine(DB_URL)
+    chunker = DrugChunker(chunk_size=1500, chunk_overlap=200, min_chunk_size=200)
 
     chunks = []
     chunk_meta = []
 
     for entry in entries:
         source = f"{entry['source']['version']} {entry['source']['volume']}"
+        drug_name = entry["name"]
         for section_key, section_text in entry["sections"].items():
             if not section_text or len(section_text.strip()) < 5:
                 continue
-            content = f"【{entry['name']} - {section_key}】\n{section_text}\n\n来源：{source}"
-            chunks.append(content)
-            chunk_meta.append({
-                "drug_id": entry["drug_id"],
-                "section": section_key,
-                "content": content,
-                "source": source,
-            })
+
+            if len(section_text) <= 1500:
+                content = f"【{drug_name} - {section_key}】\n{section_text}\n\n来源：{source}"
+                chunks.append(content)
+                chunk_meta.append({
+                    "drug_id": entry["drug_id"],
+                    "section": section_key,
+                    "content": content,
+                    "source": source,
+                })
+            else:
+                source_dict = {"version": entry["source"]["version"],
+                              "volume": entry["source"]["volume"],
+                              "page": entry["source"].get("page", "")}
+                sub_chunks = chunker._split_long_section(section_text, drug_name, section_key, source_dict)
+                for sub in sub_chunks:
+                    full_content = sub.content + f"\n\n来源：{source}"
+                    chunks.append(full_content)
+                    chunk_meta.append({
+                        "drug_id": entry["drug_id"],
+                        "section": section_key,
+                        "content": full_content,
+                        "source": source,
+                    })
 
     print(f"   {len(chunks)} chunks，向量化中...")