Prechádzať zdrojové kódy

修复java后端问题

liuchengsen 1 mesiac pred
rodič
commit
eb69a99838
1 zmenil súbory, kde vykonal 133 pridanie a 111 odobranie
  1. 133 111
      data-pipeline/docx_ingest.py

+ 133 - 111
data-pipeline/docx_ingest.py

@@ -1,6 +1,7 @@
 """
 DOCX 药典批量导入脚本
 遍历指定目录下所有 DOCX 文件 → 提取药名和 sections → 向量化 → 写入 PostgreSQL
+药品名以文件名(去 .docx)为准
 """
 import json
 import asyncio
@@ -20,7 +21,9 @@ from sqlalchemy import text
 # 配置
 # ============================================
 SOURCE_DIR = os.environ.get("DOCX_SOURCE_DIR", "/opt/2025")
-SKIP_DIRS = {"凡例", "通则", "凡例与通则", "附录", "索引", "目录"}
+
+# 跳过的目录名(非药品内容:凡例、纲要、通则说明等)
+SKIP_DIRS = {"索引", "目录", "概述"}
 
 # 从环境变量读取
 PG_HOST = os.environ.get("POSTGRES_HOST", "localhost")
@@ -38,35 +41,29 @@ BATCH_SIZE = 10  # 每批向量化的文本数
 
 # Section 头部关键词 → 标准名称映射
 SECTION_HEADERS = OrderedDict([
-    ("处方", "处方"),
-    ("制法", "制法"),
-    ("性状", "性状"),
-    ("鉴别", "鉴别"),
-    ("检查", "检查"),
-    ("浸出物", "浸出物"),
-    ("含量测定", "含量测定"),
-    ("含量", "含量测定"),
-    ("功能与主治", "功能主治"),
-    ("功能", "功能主治"),
-    ("主治", "功能主治"),
-    ("用法与用量", "用法用量"),
-    ("用法", "用法用量"),
-    ("用量", "用法用量"),
-    ("注意", "注意事项"),
-    ("注意事项", "注意事项"),
-    ("规格", "规格"),
-    ("贮藏", "贮藏"),
-    ("类别", "类别"),
-    ("制剂", "制剂"),
-    ("附注", "附注"),
+    ("处方", "处方"), ("制法", "制法"), ("性状", "性状"),
+    ("鉴别", "鉴别"), ("检查", "检查"), ("浸出物", "浸出物"),
+    ("含量测定", "含量测定"), ("含量", "含量测定"),
+    ("功能与主治", "功能主治"), ("功能", "功能主治"), ("主治", "功能主治"),
+    ("用法与用量", "用法用量"), ("用法", "用法用量"), ("用量", "用法用量"),
+    ("注意", "注意事项"), ("注意事项", "注意事项"),
+    ("规格", "规格"), ("贮藏", "贮藏"), ("类别", "类别"),
+    ("制剂", "制剂"), ("附注", "附注"),
 ])
 
-# 药典部别推断
+# 药典部别映射
 VOLUME_MAP = {"output": "一部", "output2": "二部", "output3": "三部", "output4": "四部"}
 
+# 四部子目录 → 具体分类
+CATEGORY_OVERRIDE = {
+    "药用辅料": "药用辅料",
+    "通用技术要求": "通则",
+    "指导原则": "通则",
+    "通用技术要求/指导原则": "通则",
+}
+
 
 def load_env():
-    """加载 .env 文件"""
     env_file = Path(__file__).resolve().parent.parent / ".env"
     if env_file.exists():
         with open(env_file, encoding="utf-8") as f:
@@ -77,110 +74,86 @@ def load_env():
                     os.environ.setdefault(key.strip(), val.strip())
 
 
-def normalize_name(name: str) -> str:
-    """清理药名,去掉拼音行、空格等"""
-    # 移除拼音行(包含英文字母的行)
-    lines = name.split('\n')
-    result = []
-    for line in lines:
-        line = line.strip()
-        if not line:
-            continue
-        # 跳过纯拼音行(大量英文字母)
-        if sum(1 for c in line if c.isascii() and c.isalpha()) > len(line) * 0.4:
-            continue
-        # 跳过纯数字行
-        if line.replace('.', '').replace(' ', '').isdigit():
-            continue
-        result.append(line)
-    return result[0] if result else name.strip()
+def parse_docx(filepath: str) -> dict | None:
+    """
+    解析单个 DOCX 文件为药典条目。
+    药品名 = 文件名(去掉 .docx)
+    凡例/纲要类:名称加上部别前缀
+    """
+    filename = Path(filepath).stem
+    volume = get_volume_group(filepath)
+    parent_dir = Path(filepath).parent.name
 
+    # 凡例/纲要/通则类:名称格式为"一部-凡例"以区分
+    is_meta = parent_dir in ("凡例", "纲要", "通用技术要求", "指导原则")
+    display_name = f"{volume}-{filename}" if is_meta else filename
 
-def parse_docx(filepath: str) -> dict | None:
-    """解析单个 DOCX 文件为药典条目"""
     try:
         doc = Document(filepath)
     except Exception:
         return None
 
-    drug_name = None
     pinyin = ""
     sections = OrderedDict()
-    current_section = "正文"  # 第一个 heading 之前的内容
+    current_section = "正文"
     current_text = []
 
-    # 检测并提取 section 头部的正则
-    section_pattern = re.compile(
-        r'^【(.+?)】'  # 【性状】
-        r'|^(\S{2,4})$'  # 性状, 鉴别 等单独的 section header
-    )
-
-    # 第一次 pass:收集所有段落文本
-    all_paras = []
+    # 收集所有段落
     for p in doc.paragraphs:
         text = p.text.strip()
         if not text or text.isspace():
             continue
-        all_paras.append((p.style.name, text))
 
-    for style, text in all_paras:
-        # Heading 1 = 药品名
-        if 'Heading 1' in style and drug_name is None:
-            drug_name = normalize_name(text)
+        style = p.style.name if p.style else ""
+
+        # Heading 1 通常是药品名(与文件名一致),直接跳过,不加入 sections
+        if 'Heading 1' in style:
             continue
 
-        # 检测拼音行(紧跟药名,单独一段包含字母
-        if drug_name and not pinyin and is_pinyin_line(text):
+        # 拼音行(紧跟药名后)
+        if not pinyin and is_pinyin_line(text):
             pinyin = text.strip()
             continue
 
-        # 检测 section 头部
+        # 独立的 section 标题(2-6 字短行)
         section_matched = detect_section_header(text)
         if section_matched and len(text) <= 10:
-            # 保存上一个 section
             if current_text:
                 sections[current_section] = '\n'.join(current_text).strip()
                 current_text = []
             current_section = section_matched
             continue
 
-        # 检测行内 section 头部:【性状】xxx 格式
+        # 行内 section:【性状】xxx
         m = re.match(r'^【(.+?)】\s*(.*)', text)
-        if m:
-            header = m.group(1)
-            rest = m.group(2)
-            # 看看 header 是否匹配已知 section
-            matched = detect_section_header(f"【{header}】")
-            if matched:
-                if current_text:
-                    sections[current_section] = '\n'.join(current_text).strip()
-                    current_text = []
-                current_section = matched
-                if rest:
-                    current_text.append(rest)
-                continue
+        if m and detect_section_header(m.group(1)):
+            if current_text:
+                sections[current_section] = '\n'.join(current_text).strip()
+                current_text = []
+            current_section = detect_section_header(m.group(1))
+            if m.group(2):
+                current_text.append(m.group(2))
+            continue
 
-        # 普通段落,追加到当前 section
+        # 普通段落
         current_text.append(text)
 
     # 最后一个 section
     if current_text:
         sections[current_section] = '\n'.join(current_text).strip()
 
-    if not drug_name or len(sections) == 0:
+    if len(sections) == 0:
         return None
 
-    # 推断分类和部别
-    parent_dir = Path(filepath).parent.name
-    volume = get_volume_group(filepath)
+    category = infer_category(filepath, parent_dir)
 
     return {
-        "drug_id": generate_drug_id(drug_name, filepath),
-        "name": drug_name,
+        "drug_id": generate_drug_id(display_name, filepath),
+        "name": display_name,
         "name_en": "",
         "pinyin": pinyin,
-        "category": infer_category(filepath),
-        "subcategory": parent_dir if parent_dir not in ("output", "output2", "output3", "output4") else "",
+        "category": category,
+        "subcategory": infer_subcategory(filepath),
         "sections": dict(sections),
         "source": {
             "version": "2025年版",
@@ -191,13 +164,11 @@ def parse_docx(filepath: str) -> dict | None:
 
 
 def is_pinyin_line(text: str) -> bool:
-    """判断是否为拼音行"""
     alpha_count = sum(1 for c in text if c.isascii() and c.isalpha())
     return alpha_count > len(text) * 0.3 and len(text) > 3 and len(text) < 200
 
 
 def detect_section_header(text: str) -> str | None:
-    """检测是否为 section 标题行,返回标准 section 名称"""
     clean = text.replace("【", "").replace("】", "").strip()
     if clean in SECTION_HEADERS:
         return SECTION_HEADERS[clean]
@@ -205,25 +176,47 @@ def detect_section_header(text: str) -> str | None:
 
 
 def get_volume_group(filepath: str) -> str:
-    """推断药典部别"""
     for key, vol in VOLUME_MAP.items():
         if key in filepath:
             return vol
     return "一部"
 
 
-def infer_category(filepath: str) -> str:
-    """推断药品分类"""
+def infer_category(filepath: str, parent_dir: str = "") -> str:
     path_lower = filepath.lower()
-    if "output2" in path_lower:
-        return "化学药"
-    if "output3" in path_lower:
-        return "生物制品"
+    # 凡例/纲要/通用技术要求
+    if parent_dir in ("凡例", "纲要"):
+        return "凡例与纲要"
+    if parent_dir in ("通用技术要求", "指导原则"):
+        return "通则"
+    if "output2" in path_lower: return "化学药"
+    if "output3" in path_lower: return "生物制品"
+    if "output4" in path_lower: return _infer_category_vol4(filepath)
     return "中药"
 
 
+def _infer_category_vol4(filepath: str) -> str:
+    """四部特殊处理"""
+    for key, cat in CATEGORY_OVERRIDE.items():
+        if key in filepath:
+            return cat
+    return "通则"
+
+
+def infer_subcategory(filepath: str) -> str:
+    """从目录层级推断子分类"""
+    parts = Path(filepath).parts
+    # 找到 output/output2/output3/output4 之后的子目录名
+    for i, p in enumerate(parts):
+        if p in ("output", "output2", "output3", "output4"):
+            if i + 1 < len(parts):
+                sub = parts[i + 1]
+                # 排除数字编号前缀
+                return sub
+    return ""
+
+
 def generate_drug_id(name: str, filepath: str) -> str:
-    """生成 drug_id"""
     vol = get_volume_group(filepath)
     prefix = {"一部": "Z", "二部": "H", "三部": "S", "四部": "T"}.get(vol, "Z")
     hash_suffix = hashlib.md5(name.encode()).hexdigest()[:6].upper()
@@ -234,13 +227,10 @@ def find_docx_files(root_dir: str) -> list[str]:
     """扫描所有 DOCX 文件,跳过非药品目录"""
     files = []
     for dirpath, dirnames, filenames in os.walk(root_dir):
-        # 跳过非药品目录
         dir_basename = os.path.basename(dirpath)
         if dir_basename in SKIP_DIRS:
             dirnames.clear()
             continue
-
-        # 跳过临时文件
         for f in sorted(filenames):
             if f.startswith('~') or f.startswith('.'):
                 continue
@@ -250,7 +240,6 @@ def find_docx_files(root_dir: str) -> list[str]:
 
 
 async def get_embeddings(texts: list[str], text_type: str = "document") -> list[list[float]]:
-    """调用 Qwen Embedding API"""
     async with httpx.AsyncClient(timeout=60) as client:
         resp = await client.post(
             EMBEDDING_URL,
@@ -272,7 +261,6 @@ async def get_embeddings(texts: list[str], text_type: str = "document") -> list[
 
 async def ingest_docx_entries(entries: list[dict], engine, start_idx: int = 0):
     """批量向量化 + 入库"""
-    # 构建 chunk 文本
     chunks = []
     chunk_meta = []
 
@@ -292,7 +280,6 @@ async def ingest_docx_entries(entries: list[dict], engine, start_idx: int = 0):
 
     print(f"  ✂️  {len(chunks)} 个 chunks,向量化中...")
 
-    # 分批向量化
     all_vectors = []
     for i in range(0, len(chunks), BATCH_SIZE):
         batch = chunks[i:i + BATCH_SIZE]
@@ -343,6 +330,9 @@ async def ingest_docx_entries(entries: list[dict], engine, start_idx: int = 0):
                     ON CONFLICT (drug_id) DO UPDATE SET
                         sections = EXCLUDED.sections,
                         category = EXCLUDED.category,
+                        subcategory = EXCLUDED.subcategory,
+                        source_version = EXCLUDED.source_version,
+                        source_volume = EXCLUDED.source_volume,
                         updated_at = NOW()
                 """),
                 {
@@ -363,6 +353,27 @@ async def ingest_docx_entries(entries: list[dict], engine, start_idx: int = 0):
     return drug_count, chunk_count
 
 
+async def dedup_old_versions(engine):
+    """覆盖性新增:同名药品如果存在非 2025 版的旧记录,标记为 inactive"""
+    async with engine.begin() as conn:
+        result = await conn.execute(
+            text("""
+                UPDATE drugs
+                SET is_active = FALSE, updated_at = NOW()
+                WHERE is_active = TRUE
+                  AND source_version != '2025年版'
+                  AND name IN (
+                      SELECT name FROM drugs
+                      WHERE source_version = '2025年版' AND is_active = TRUE
+                  )
+            """)
+        )
+        # asyncpg 不支持 rowcount,用另一种方式
+        count = getattr(result, 'rowcount', 0)
+        if count > 0:
+            print(f"  🧹 去重: 标记 {count} 个旧版药品为 inactive")
+
+
 async def main():
     load_env()
     global QWEN_API_KEY
@@ -377,30 +388,38 @@ async def main():
     print(f"   发现 {len(files)} 个 DOCX 文件")
     print("=" * 60)
 
-    # 解析所有文件
+    # 解析并统计
     entries = []
-    parse_errors = 0
+    skipped = 0
+    by_category = {}
     for i, fp in enumerate(files):
         entry = parse_docx(fp)
         if entry:
             entries.append(entry)
-
-        if (i + 1) % 500 == 0:
+            cat = entry["category"]
+            by_category[cat] = by_category.get(cat, 0) + 1
+        else:
+            skipped += 1
+        if (i + 1) % 1000 == 0:
             print(f"   解析进度: {i + 1}/{len(files)} (有效: {len(entries)})")
 
-    print(f"\n📦 解析完成: {len(entries)} 个有效药品条目 (跳过 {len(files) - len(entries)} 个)")
+    print(f"\n📦 解析完成: {len(entries)} 个有效条目 (跳过 {skipped} 个)")
+    print("   分类分布:")
+    for cat, cnt in sorted(by_category.items(), key=lambda x: -x[1]):
+        print(f"     {cat}: {cnt}")
 
     if not entries:
         print("❌ 没有有效数据")
         return
 
-    # 连接数据库,分批入库
+    # 连接数据库
     engine = create_async_engine(DB_URL)
+
+    # 分批入库
     total_drugs = 0
     total_chunks = 0
-
-    # 分批处理(每批 50 个药品,避免 Qwen API 超时)
     BATCH = 50
+
     for i in range(0, len(entries), BATCH):
         batch = entries[i:i + BATCH]
         batch_no = i // BATCH + 1
@@ -413,14 +432,17 @@ async def main():
             print(f"   ✅ 入库: {dc} 药品, {cc} chunks")
         except Exception as e:
             print(f"   ❌ 批次失败: {e}")
-            # 打印第一个药品名用于调试
             if batch:
                 print(f"      首个药品: {batch[0]['name']}")
 
+    # 去重:标记同名旧版药品为 inactive
+    print("\n🧹 清理旧版重复数据...")
+    await dedup_old_versions(engine)
+
     await engine.dispose()
 
     print("\n" + "=" * 60)
-    print(f"🎉 全部完成!药品 {total_drugs} 个, chunks {total_chunks} 条")
+    print(f"🎉 全部完成!新增药品 {total_drugs} 个, chunks {total_chunks} 条")
     print("=" * 60)