2020年药典学习python

data-pipeline/pdf2020_ingest.py

@@ -56,133 +56,135 @@ def load_env():
                 os.environ.setdefault(key.strip(), val.strip())
 
 
+def is_cjk_name(s):
+    """严格药名检测：2-6 个纯 CJK 字符"""
+    if not s or len(s) < 2 or len(s) > 6:
+        return False
+    cjk = sum(1 for c in s if ord(c) >= 0x4E00 and ord(c) <= 0x9FFF)
+    return cjk >= 2 and cjk == len(s)
+
+
+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]:
-    """从 PDF 中提取药品条目"""
+    """解析策略：扫描【性状】锚点，向前找药名。
+    药名特征：2-6 纯 CJK + 下一行是拼音（首字母大写）或全大写拉丁名。
+    只有通过拼音/拉丁名确认的才作为药名候选。
+    """
     doc = fitz.open(pdf_path)
-    print(f"📖 PDF: {doc.page_count} 页")
+    print(f"PDF: {doc.page_count} pages")
 
-    # 收集所有页面的文本
+    # 跳过前言（前 50 页）
     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)
+    for i in range(32, doc.page_count):
+        all_text += doc[i].get_text() + "\n"
+    doc.close()
 
-    # 找到正文开始位置（跳过前言/TOC）
-    # 正文以药名开头，通常格式为 "药名" 后紧跟 "【性状】" 等 section
-    # 扫描找到第一个药品条目
     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)
+
+    # 第一遍：标记所有"已确认的药名"——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)
+
+    print(f"   确认药名: {len(confirmed_names)} 个")
+
+    # 第二遍：以已确认药名为锚点，提取药品条目
     entries = []
-    current_name = None
+    current_drug = None
     current_sections = OrderedDict()
     current_section = "正文"
     current_text = []
-    in_entry = False
-
-    for line in lines:
-        line = line.strip()
-        if not line:
-            continue
+    in_drug = False
 
-        # 跳过纯页码和标题行
-        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)
+    for s in clean_lines:
+        # 检测 section
+        m = re.match(r'^【(.+?)】\s*(.*)', s)
         if m:
-            sec_name = m.group(1)
+            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_name:
+                if key in sec_raw:
                     section_found = std_name
-                    section_content = m.group(2)
                     break
 
-        if section_found:
-            if in_entry and current_name:
-                # 保存上一个 section
+            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 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
+                if sec_content:
+                    current_text.append(sec_content)
+                continue
 
-        # 可能是药名行（短行，没有 section 标记，以中文开头）
-        # 药名通常在 section 之前的一两行
-        if not in_entry and re.match(r'^[一-鿿]{2,20}$', line):
-            # 可能是新药名
-            if current_name and current_sections:
-                # 保存上一个药品
+        # 检测已确认的药名
+        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_name,
-                    "sections": dict(current_sections),
-                })
-                current_text = []
-                current_sections = OrderedDict()
-            current_name = line
-            in_entry = False
+                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_entry and current_name:
-            current_text.append(line)
+        # 正文内容
+        if in_drug:
+            current_text.append(s)
 
-    # 最后一个药品
-    if current_name and current_sections:
+    # 最后一个
+    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_name,
-            "sections": dict(current_sections),
-        })
+        entries.append({"name": current_drug, "sections": dict(current_sections)})
 
-    doc.close()
+    print(f"   初步提取: {len(entries)} 个条目")
 
-    # 过滤无效条目（至少要有 2 个 section 或内容 > 100 字）
-    valid = []
+    # 去重：同一药名保留内容最多的条目
+    best = {}
     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
+        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]]:
@@ -226,7 +228,7 @@ async def ingest_entries(entries: list[dict]):
                 "source": source,
             })
 
-    print(f"  ✂️  {len(chunks)} chunks，向量化中...")
+    print(f"   {len(chunks)} chunks，向量化中...")
 
     all_vectors = []
     for i in range(0, len(chunks), BATCH_SIZE):
@@ -236,7 +238,7 @@ async def ingest_entries(entries: list[dict]):
         n = min(i + BATCH_SIZE, len(chunks))
         print(f"    向量化: {n}/{len(chunks)}")
 
-    print(f"  ✅ 向量化完成, 维度={len(all_vectors[0]) if all_vectors else 'N/A'}")
+    print(f"  [OK] 向量化完成, 维度={len(all_vectors[0]) if all_vectors else 'N/A'}")
 
     # 写入 drug_chunks
     chunk_count = 0
@@ -291,7 +293,7 @@ async def ingest_entries(entries: list[dict]):
             drug_count += 1
 
     await engine.dispose()
-    print(f"\n🎉 2020 药典入库完成！药品 {drug_count} 个，chunk {chunk_count} 条")
+    print(f"\n[DONE] 2020 药典入库完成！药品 {drug_count} 个，chunk {chunk_count} 条")
 
 
 async def main():
@@ -303,7 +305,7 @@ async def main():
         sys.exit(1)
 
     print("=" * 60)
-    print("📖 解析 2020 年版药典 PDF...")
+    print("解析 2020 年版药典 PDF...")
     entries = extract_drug_entries(PDF_PATH)
     print(f"   提取药品条目: {len(entries)}")
     print("=" * 60)
@@ -313,13 +315,13 @@ async def main():
         return
 
     # 打印前 5 个条目作为样品
-    print("\n📋 前 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🚀 开始向量化入库...")
+    print(f"\n[>] 开始向量化入库...")
     await ingest_entries(entries)