|
|
@@ -0,0 +1,327 @@
|
|
|
+"""
|
|
|
+2020 年版药典 PDF 导入脚本
|
|
|
+解析 pharmacopoeia_2020_volume1_toc.pdf → 提取药品条目 → 向量化 → 写入 PostgreSQL
|
|
|
+"""
|
|
|
+import json, os, sys, re, hashlib, asyncio
|
|
|
+from pathlib import Path
|
|
|
+from collections import OrderedDict
|
|
|
+
|
|
|
+import httpx
|
|
|
+import fitz # pymupdf
|
|
|
+from sqlalchemy.ext.asyncio import create_async_engine
|
|
|
+from sqlalchemy import text
|
|
|
+
|
|
|
+# ============================================
|
|
|
+# 配置
|
|
|
+# ============================================
|
|
|
+PDF_PATH = os.environ.get("PDF_2020_PATH", os.path.join(
|
|
|
+ os.path.dirname(__file__), "..", "data", "pharmacopoeia_2020_volume1_toc.pdf"))
|
|
|
+
|
|
|
+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
|
|
|
+
|
|
|
+# 2020 药典 section 标题识别
|
|
|
+SECTION_HEADERS = OrderedDict([
|
|
|
+ ("性状", "性状"), ("鉴别", "鉴别"), ("检查", "检查"),
|
|
|
+ ("含量测定", "含量测定"), ("浸出物", "浸出物"),
|
|
|
+ ("性味与归经", "性味与归经"), ("性味", "性味与归经"),
|
|
|
+ ("功能与主治", "功能主治"), ("功能", "功能主治"), ("主治", "功能主治"),
|
|
|
+ ("用法与用量", "用法用量"), ("用法", "用法用量"), ("用量", "用法用量"),
|
|
|
+ ("注意", "注意事项"), ("注意事项", "注意事项"),
|
|
|
+ ("规格", "规格"), ("贮藏", "贮藏"), ("类别", "类别"),
|
|
|
+ ("制剂", "制剂"), ("附注", "附注"), ("禁忌", "禁忌"),
|
|
|
+ ("不良反应", "不良反应"),
|
|
|
+ ("处方", "处方"), ("制法", "制法"),
|
|
|
+ ("包装", "包装"), ("有效期", "有效期"),
|
|
|
+ ("执行标准", "执行标准"), ("批准文号", "批准文号"),
|
|
|
+])
|
|
|
+
|
|
|
+
|
|
|
+def load_env():
|
|
|
+ env_file = Path(__file__).resolve().parent.parent / ".env"
|
|
|
+ if env_file.exists():
|
|
|
+ for line in open(env_file, encoding="utf-8"):
|
|
|
+ line = line.strip()
|
|
|
+ if line and not line.startswith("#") and "=" in line:
|
|
|
+ key, _, val = line.partition("=")
|
|
|
+ os.environ.setdefault(key.strip(), val.strip())
|
|
|
+
|
|
|
+
|
|
|
+def extract_drug_entries(pdf_path: str) -> list[dict]:
|
|
|
+ """从 PDF 中提取药品条目"""
|
|
|
+ doc = fitz.open(pdf_path)
|
|
|
+ print(f"📖 PDF: {doc.page_count} 页")
|
|
|
+
|
|
|
+ # 收集所有页面的文本
|
|
|
+ all_text = ""
|
|
|
+ for i in range(doc.page_count):
|
|
|
+ text = doc[i].get_text()
|
|
|
+ all_text += text + "\n"
|
|
|
+
|
|
|
+ # 清理：去掉页眉页码、多余空白
|
|
|
+ all_text = re.sub(r'\n{3,}', '\n\n', all_text)
|
|
|
+ all_text = re.sub(r'^\s*\d+\s*$', '', all_text, flags=re.MULTILINE)
|
|
|
+
|
|
|
+ # 找到正文开始位置（跳过前言/TOC）
|
|
|
+ # 正文以药名开头，通常格式为 "药名" 后紧跟 "【性状】" 等 section
|
|
|
+ # 扫描找到第一个药品条目
|
|
|
+ lines = all_text.split('\n')
|
|
|
+
|
|
|
+ # 策略：扫描【性状】标记，向前找药名
|
|
|
+ entries = []
|
|
|
+ current_name = None
|
|
|
+ current_sections = OrderedDict()
|
|
|
+ current_section = "正文"
|
|
|
+ current_text = []
|
|
|
+ in_entry = False
|
|
|
+
|
|
|
+ for line in lines:
|
|
|
+ line = line.strip()
|
|
|
+ if not line:
|
|
|
+ continue
|
|
|
+
|
|
|
+ # 跳过纯页码和标题行
|
|
|
+ if re.match(r'^\d{1,4}$', line):
|
|
|
+ continue
|
|
|
+ if line.startswith("中国药典") or line.startswith("ISBN"):
|
|
|
+ continue
|
|
|
+ if "图书在版" in line:
|
|
|
+ continue
|
|
|
+
|
|
|
+ # 检测 section 标题
|
|
|
+ section_found = None
|
|
|
+ section_content = ""
|
|
|
+
|
|
|
+ # 匹配 【xxx】 格式
|
|
|
+ m = re.match(r'^【(.+?)】\s*(.*)', line)
|
|
|
+ if m:
|
|
|
+ sec_name = m.group(1)
|
|
|
+ for key, std_name in SECTION_HEADERS.items():
|
|
|
+ if key in sec_name:
|
|
|
+ section_found = std_name
|
|
|
+ section_content = m.group(2)
|
|
|
+ break
|
|
|
+
|
|
|
+ if section_found:
|
|
|
+ if in_entry and current_name:
|
|
|
+ # 保存上一个 section
|
|
|
+ if current_text:
|
|
|
+ current_sections[current_section] = '\n'.join(current_text).strip()
|
|
|
+ current_text = []
|
|
|
+ current_section = section_found
|
|
|
+ if section_content:
|
|
|
+ current_text.append(section_content)
|
|
|
+ elif not in_entry and current_name:
|
|
|
+ # 第一个 section，标志着药品条目开始
|
|
|
+ in_entry = True
|
|
|
+ current_sections = OrderedDict()
|
|
|
+ current_section = section_found
|
|
|
+ current_text = [section_content] if section_content else []
|
|
|
+ continue
|
|
|
+
|
|
|
+ # 可能是药名行（短行，没有 section 标记，以中文开头）
|
|
|
+ # 药名通常在 section 之前的一两行
|
|
|
+ if not in_entry and re.match(r'^[一-鿿]{2,20}$', line):
|
|
|
+ # 可能是新药名
|
|
|
+ if current_name and current_sections:
|
|
|
+ # 保存上一个药品
|
|
|
+ if current_text:
|
|
|
+ current_sections[current_section] = '\n'.join(current_text).strip()
|
|
|
+ entries.append({
|
|
|
+ "name": current_name,
|
|
|
+ "sections": dict(current_sections),
|
|
|
+ })
|
|
|
+ current_text = []
|
|
|
+ current_sections = OrderedDict()
|
|
|
+ current_name = line
|
|
|
+ in_entry = False
|
|
|
+ current_section = "正文"
|
|
|
+ current_text = []
|
|
|
+ continue
|
|
|
+
|
|
|
+ # 正文内容行
|
|
|
+ if in_entry and current_name:
|
|
|
+ current_text.append(line)
|
|
|
+
|
|
|
+ # 最后一个药品
|
|
|
+ if current_name and current_sections:
|
|
|
+ if current_text:
|
|
|
+ current_sections[current_section] = '\n'.join(current_text).strip()
|
|
|
+ entries.append({
|
|
|
+ "name": current_name,
|
|
|
+ "sections": dict(current_sections),
|
|
|
+ })
|
|
|
+
|
|
|
+ doc.close()
|
|
|
+
|
|
|
+ # 过滤无效条目（至少要有 2 个 section 或内容 > 100 字）
|
|
|
+ valid = []
|
|
|
+ for e in entries:
|
|
|
+ content_len = sum(len(v) for v in e["sections"].values())
|
|
|
+ if len(e["sections"]) >= 2 or content_len > 100:
|
|
|
+ # 生成 drug_id
|
|
|
+ hash_suffix = hashlib.md5(e["name"].encode()).hexdigest()[:6].upper()
|
|
|
+ e["drug_id"] = f"Z2020-{hash_suffix}"
|
|
|
+ e["category"] = "中药"
|
|
|
+ e["subcategory"] = ""
|
|
|
+ e["name_en"] = ""
|
|
|
+ e["pinyin"] = ""
|
|
|
+ e["source"] = {
|
|
|
+ "version": "2020年版",
|
|
|
+ "volume": "一部",
|
|
|
+ "page": "",
|
|
|
+ }
|
|
|
+ valid.append(e)
|
|
|
+
|
|
|
+ return valid
|
|
|
+
|
|
|
+
|
|
|
+async def get_embeddings(texts: list[str], text_type: str = "document") -> list[list[float]]:
|
|
|
+ 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_entries(entries: list[dict]):
|
|
|
+ """向量化 + 入库"""
|
|
|
+ engine = create_async_engine(DB_URL)
|
|
|
+
|
|
|
+ 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": idx,
|
|
|
+ "embedding": json.dumps(vec),
|
|
|
+ "vec": vec_str,
|
|
|
+ },
|
|
|
+ )
|
|
|
+ chunk_count += 1
|
|
|
+
|
|
|
+ # 写入 drugs 表
|
|
|
+ drug_count = 0
|
|
|
+ async with engine.begin() as conn:
|
|
|
+ for entry in entries:
|
|
|
+ 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,
|
|
|
+ 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
|
|
|
+
|
|
|
+ await engine.dispose()
|
|
|
+ print(f"\n🎉 2020 药典入库完成！药品 {drug_count} 个，chunk {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")
|
|
|
+ sys.exit(1)
|
|
|
+
|
|
|
+ print("=" * 60)
|
|
|
+ print("📖 解析 2020 年版药典 PDF...")
|
|
|
+ entries = extract_drug_entries(PDF_PATH)
|
|
|
+ print(f" 提取药品条目: {len(entries)}")
|
|
|
+ print("=" * 60)
|
|
|
+
|
|
|
+ if not entries:
|
|
|
+ print("❌ 未提取到有效条目，请检查 PDF 解析逻辑")
|
|
|
+ return
|
|
|
+
|
|
|
+ # 打印前 5 个条目作为样品
|
|
|
+ print("\n📋 前 5 个条目预览:")
|
|
|
+ for e in entries[:5]:
|
|
|
+ secs = list(e["sections"].keys())
|
|
|
+ total_len = sum(len(v) for v in e["sections"].values())
|
|
|
+ print(f" {e['name']} | {len(secs)} sections ({total_len} 字) | {secs}")
|
|
|
+
|
|
|
+ print(f"\n🚀 开始向量化入库...")
|
|
|
+ await ingest_entries(entries)
|
|
|
+
|
|
|
+
|
|
|
+if __name__ == "__main__":
|
|
|
+ asyncio.run(main())
|
liuchengsen