""" 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 from processors.chunker import DrugChunker # ============================================ # 配置 # ============================================ 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 is_cjk_name(s): """药名检测：2-10 字符，CJK 占 70% 以上，允许破折号/连接符/括号""" if not s or len(s) < 2 or len(s) > 10: return False # Skip pure digits/punctuation if re.match(r'^[\d\s\.\,\;\:\!\?\(\)\[\]\{\}<>@#$%^&*+=|\\/\"\']+$', s): return False cjk = sum(1 for c in s if ord(c) >= 0x4E00 and ord(c) <= 0x9FFF) return cjk >= 2 and cjk >= len(s) * 0.7 def is_pinyin_line(s): """拼音行：纯拉丁字母空格，首字母大写""" if not s or len(s) < 3 or len(s) > 80: return False alpha_spaces = sum(1 for c in s if c.isalpha() or c.isspace()) return alpha_spaces == len(s) and any(c.isalpha() for c in s) and s[0].isupper() def extract_drug_entries(pdf_path: str) -> list[dict]: """解析策略：扫描【性状】锚点，向前找药名。 药名特征：2-6 纯 CJK + 下一行是拼音（首字母大写）或全大写拉丁名。 只有通过拼音/拉丁名确认的才作为药名候选。 """ doc = fitz.open(pdf_path) print(f"PDF: {doc.page_count} pages") # 跳过前言（前 50 页） all_text = "" for i in range(32, doc.page_count): all_text += doc[i].get_text() + "\n" doc.close() lines = all_text.split('\n') clean_lines = [] for line in lines: s = line.strip() if not s: continue if s.startswith("中国药典") and len(s) < 30: continue if "ISBN" in s or "CIP" in s: continue if re.match(r'^[・•·]\s*\d+\s*[・•·]$', s): continue clean_lines.append(s) # 第一遍：标记所有"已确认的药名" # 方式1：CJK + 下一行是拼音或拉丁名（药材和饮片） # 方式2：CJK + 下一行是【处方】或【制法】（成方制剂，无拼音） confirmed_names = set() for i in range(len(clean_lines) - 2): s = clean_lines[i] nxt = clean_lines[i + 1] if is_cjk_name(s): if is_pinyin_line(nxt) or (nxt.isupper() and len(nxt) > 5): confirmed_names.add(s) elif nxt.startswith('【处方】') or nxt.startswith('【制法】'): confirmed_names.add(s) print(f" 确认药名: {len(confirmed_names)} 个") # 第二遍：以已确认药名为锚点，提取药品条目 entries = [] current_drug = None current_sections = OrderedDict() current_section = "正文" current_text = [] in_drug = False for s in clean_lines: # 检测 section m = re.match(r'^【(.+?)】\s*(.*)', s) if m: sec_raw = m.group(1) sec_content = m.group(2) section_found = None for key, std_name in SECTION_HEADERS.items(): if key in sec_raw: section_found = std_name break if section_found and in_drug: if current_text: current_sections[current_section] = '\n'.join(current_text).strip() current_text = [] current_section = section_found if sec_content: current_text.append(sec_content) continue # 检测已确认的药名 if s in confirmed_names: # 保存上一个药品 if in_drug and current_drug and current_sections: if current_text: current_sections[current_section] = '\n'.join(current_text).strip() entries.append({"name": current_drug, "sections": dict(current_sections)}) # 开始新药品 current_drug = s in_drug = True current_sections = OrderedDict() current_section = "正文" current_text = [] # 描述段落（药名和【性状】之间的文本）归入"正文" continue # 跳过拼音行和拉丁名行 if in_drug and not current_sections and (is_pinyin_line(s) or (s.isupper() and len(s) > 5)): continue # 正文内容 if in_drug: current_text.append(s) # 最后一个 if in_drug and current_drug and current_sections: if current_text: current_sections[current_section] = '\n'.join(current_text).strip() entries.append({"name": current_drug, "sections": dict(current_sections)}) print(f" 初步提取: {len(entries)} 个条目") # 去重：同一药名保留内容最多的条目 best = {} for e in entries: content_len = sum(len(v) for v in e["sections"].values()) if len(e["sections"]) >= 1 and content_len > 50: name = e["name"] if name not in best or content_len > sum(len(v) for v in best[name]["sections"].values()): hash_suffix = hashlib.md5(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": ""} best[name] = e return list(best.values()) 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]): """向量化 + 入库。长 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 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，向量化中...") 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" [OK] 向量化完成, 维度={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[DONE] 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())