feat: docx药典批量导入脚本 + 修复ingest.py vec列

data-pipeline/docx_ingest.py

@@ -0,0 +1,428 @@
+"""
+DOCX 药典批量导入脚本
+遍历指定目录下所有 DOCX 文件 → 提取药名和 sections → 向量化 → 写入 PostgreSQL
+"""
+import json
+import asyncio
+import os
+import sys
+import re
+import hashlib
+from pathlib import Path
+from collections import OrderedDict
+
+import httpx
+from docx import Document
+from sqlalchemy.ext.asyncio import create_async_engine
+from sqlalchemy import text
+
+# ============================================
+# 配置
+# ============================================
+SOURCE_DIR = os.environ.get("DOCX_SOURCE_DIR", "/opt/2025")
+SKIP_DIRS = {"凡例", "通则", "凡例与通则", "附录", "索引", "目录"}
+
+# 从环境变量读取
+PG_HOST = os.environ.get("POSTGRES_HOST", "localhost")
+PG_PORT = os.environ.get("POSTGRES_PORT", "5432")
+PG_DB = os.environ.get("POSTGRES_DB", "pharmacopoeia")
+PG_USER = os.environ.get("POSTGRES_USER", "postgres")
+PG_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "postgres")
+DB_URL = f"postgresql+asyncpg://{PG_USER}:{PG_PASSWORD}@{PG_HOST}:{PG_PORT}/{PG_DB}"
+
+QWEN_API_KEY = os.environ.get("QWEN_API_KEY", "")
+EMBEDDING_URL = "https://dashscope.aliyuncs.com/api/v1/services/embeddings/text-embedding/text-embedding"
+EMBEDDING_MODEL = "text-embedding-v3"
+
+BATCH_SIZE = 10  # 每批向量化的文本数
+
+# Section 头部关键词 → 标准名称映射
+SECTION_HEADERS = OrderedDict([
+    ("处方", "处方"),
+    ("制法", "制法"),
+    ("性状", "性状"),
+    ("鉴别", "鉴别"),
+    ("检查", "检查"),
+    ("浸出物", "浸出物"),
+    ("含量测定", "含量测定"),
+    ("含量", "含量测定"),
+    ("功能与主治", "功能主治"),
+    ("功能", "功能主治"),
+    ("主治", "功能主治"),
+    ("用法与用量", "用法用量"),
+    ("用法", "用法用量"),
+    ("用量", "用法用量"),
+    ("注意", "注意事项"),
+    ("注意事项", "注意事项"),
+    ("规格", "规格"),
+    ("贮藏", "贮藏"),
+    ("类别", "类别"),
+    ("制剂", "制剂"),
+    ("附注", "附注"),
+])
+
+# 药典部别推断
+VOLUME_MAP = {"output": "一部", "output2": "二部", "output3": "三部", "output4": "四部"}
+
+
+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:
+            for line in f:
+                line = line.strip()
+                if line and not line.startswith("#") and "=" in line:
+                    key, _, val = line.partition("=")
+                    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 文件为药典条目"""
+    try:
+        doc = Document(filepath)
+    except Exception:
+        return None
+
+    drug_name = None
+    pinyin = ""
+    sections = OrderedDict()
+    current_section = "正文"  # 第一个 heading 之前的内容
+    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)
+            continue
+
+        # 检测拼音行（紧跟药名之后，单独一段包含字母）
+        if drug_name and not pinyin and is_pinyin_line(text):
+            pinyin = text.strip()
+            continue
+
+        # 检测 section 头部
+        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 格式
+        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
+
+        # 普通段落，追加到当前 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:
+        return None
+
+    # 推断分类和部别
+    parent_dir = Path(filepath).parent.name
+    volume = get_volume_group(filepath)
+
+    return {
+        "drug_id": generate_drug_id(drug_name, filepath),
+        "name": drug_name,
+        "name_en": "",
+        "pinyin": pinyin,
+        "category": infer_category(filepath),
+        "subcategory": parent_dir if parent_dir not in ("output", "output2", "output3", "output4") else "",
+        "sections": dict(sections),
+        "source": {
+            "version": "2025年版",
+            "volume": volume,
+            "page": "",
+        },
+    }
+
+
+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]
+    return 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:
+    """推断药品分类"""
+    path_lower = filepath.lower()
+    if "output2" in path_lower:
+        return "化学药"
+    if "output3" in path_lower:
+        return "生物制品"
+    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()
+    return f"{prefix}2025-{hash_suffix}"
+
+
+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
+            if f.endswith('.docx'):
+                files.append(os.path.join(dirpath, f))
+    return files
+
+
+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,
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {QWEN_API_KEY}",
+            },
+            json={
+                "model": EMBEDDING_MODEL,
+                "input": {"texts": texts},
+                "parameters": {"text_type": text_type},
+            },
+        )
+        data = resp.json()
+    if data.get("code") and data.get("code") != "" and data.get("code") is not None:
+        raise RuntimeError(f"Embedding error: {data.get('message', data)}")
+    return [item["embedding"] for item in data["output"]["embeddings"]]
+
+
+async def ingest_docx_entries(entries: list[dict], engine, start_idx: int = 0):
+    """批量向量化 + 入库"""
+    # 构建 chunk 文本
+    chunks = []
+    chunk_meta = []
+
+    for entry in entries:
+        source = f"{entry['source']['version']} {entry['source']['volume']}"
+        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,
+            })
+
+    print(f"  ✂️  {len(chunks)} 个 chunks，向量化中...")
+
+    # 分批向量化
+    all_vectors = []
+    for i in range(0, len(chunks), BATCH_SIZE):
+        batch = chunks[i:i + BATCH_SIZE]
+        vecs = await get_embeddings(batch, text_type="document")
+        all_vectors.extend(vecs)
+        n = min(i + BATCH_SIZE, len(chunks))
+        print(f"    向量化: {n}/{len(chunks)}")
+
+    print(f"  ✅ 向量化完成, 维度={len(all_vectors[0]) if all_vectors else 'N/A'}")
+
+    # 写入 drug_chunks
+    chunk_count = 0
+    async with engine.begin() as conn:
+        for idx, (meta, vec) in enumerate(zip(chunk_meta, all_vectors)):
+            vec_str = f"[{','.join(str(v) for v in vec)}]"
+            await conn.execute(
+                text("""
+                    INSERT INTO drug_chunks (drug_id, section, content, source, chunk_index, embedding, vec)
+                    VALUES (:drug_id, :section, :content, :source, :chunk_index, :embedding, :vec)
+                    ON CONFLICT DO NOTHING
+                """),
+                {
+                    "drug_id": meta["drug_id"],
+                    "section": meta["section"],
+                    "content": meta["content"],
+                    "source": meta["source"],
+                    "chunk_index": start_idx + idx,
+                    "embedding": json.dumps(vec),
+                    "vec": vec_str,
+                },
+            )
+            chunk_count += 1
+
+    # 写入 drugs 表
+    drug_count = 0
+    async with engine.begin() as conn:
+        seen = set()
+        for entry in entries:
+            if entry["drug_id"] in seen:
+                continue
+            seen.add(entry["drug_id"])
+            await conn.execute(
+                text("""
+                    INSERT INTO drugs (drug_id, name, name_en, pinyin, category, subcategory,
+                                       sections, source_version, source_volume, source_page, is_active)
+                    VALUES (:drug_id, :name, :name_en, :pinyin, :category, :subcategory,
+                            :sections, :source_version, :source_volume, :source_page, TRUE)
+                    ON CONFLICT (drug_id) DO UPDATE SET
+                        sections = EXCLUDED.sections,
+                        category = EXCLUDED.category,
+                        updated_at = NOW()
+                """),
+                {
+                    "drug_id": entry["drug_id"],
+                    "name": entry["name"],
+                    "name_en": entry.get("name_en", ""),
+                    "pinyin": entry.get("pinyin", ""),
+                    "category": entry.get("category", ""),
+                    "subcategory": entry.get("subcategory", ""),
+                    "sections": json.dumps(entry.get("sections", {}), ensure_ascii=False),
+                    "source_version": entry["source"]["version"],
+                    "source_volume": entry["source"]["volume"],
+                    "source_page": entry["source"].get("page", ""),
+                },
+            )
+            drug_count += 1
+
+    return drug_count, chunk_count
+
+
+async def main():
+    load_env()
+    global QWEN_API_KEY
+    QWEN_API_KEY = os.environ.get("QWEN_API_KEY", "")
+    if not QWEN_API_KEY:
+        print("❌ 未设置 QWEN_API_KEY，请在 .env 文件中配置")
+        sys.exit(1)
+
+    print("=" * 60)
+    print("📂 扫描 DOCX 文件...")
+    files = find_docx_files(SOURCE_DIR)
+    print(f"   发现 {len(files)} 个 DOCX 文件")
+    print("=" * 60)
+
+    # 解析所有文件
+    entries = []
+    parse_errors = 0
+    for i, fp in enumerate(files):
+        entry = parse_docx(fp)
+        if entry:
+            entries.append(entry)
+
+        if (i + 1) % 500 == 0:
+            print(f"   解析进度: {i + 1}/{len(files)} (有效: {len(entries)})")
+
+    print(f"\n📦 解析完成: {len(entries)} 个有效药品条目 (跳过 {len(files) - len(entries)} 个)")
+
+    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
+        total_batches = (len(entries) + BATCH - 1) // BATCH
+        print(f"\n🚀 第 {batch_no}/{total_batches} 批 ({len(batch)} 个药品) ...")
+        try:
+            dc, cc = await ingest_docx_entries(batch, engine, start_idx=total_chunks)
+            total_drugs += dc
+            total_chunks += cc
+            print(f"   ✅ 入库: {dc} 药品, {cc} chunks")
+        except Exception as e:
+            print(f"   ❌ 批次失败: {e}")
+            # 打印第一个药品名用于调试
+            if batch:
+                print(f"      首个药品: {batch[0]['name']}")
+
+    await engine.dispose()
+
+    print("\n" + "=" * 60)
+    print(f"🎉 全部完成！药品 {total_drugs} 个, chunks {total_chunks} 条")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())