"""遍历 _001__clawler/中药/ 下的中药详情 txt，调用 LLM 抽取实体和关系，一药一个 json。
支持断点续抽：已经存在对应 json 的中药直接跳过，中断后重跑即可继续。
输出：
_002__extract_information/中药图谱/<中药名>.json 成功（合法的空结果也算成功）
_002__extract_information/中药图谱/_failed.log 失败清单，每行「中药名<TAB>原因」
_002__extract_information/alpaca_herb_dataset.jsonl 顺便追加的 Alpaca 微调样本
数据集是抽取时顺手追加的，所以只覆盖「本次抽的」那些药材。已经抽过的药材会被
断点续抽跳过，不会再写进数据集 —— 要补齐全量或去重，跑 02_build_alpaca_dataset.py
从 中药图谱/*.json 重建。
单条 LLM 调用实测 20 秒~4 分钟（抽出来的实体关系多，输出 token 就多），所以默认 10 线程并发。
1472 条串行大约 90 小时，10 线程约 9~15 小时。想快就把 --workers 提上去，注意接口限流。
用法（在项目根目录执行）：
python -X utf8 -m _002__extract_information.01_extract_herb_graph --limit 5 # 先抽 5 条看效果
python -X utf8 -m _002__extract_information.01_extract_herb_graph # 全量，10 线程
python -X utf8 -m _002__extract_information.01_extract_herb_graph --workers 20
python -X utf8 -m _002__extract_information.01_extract_herb_graph --retry-failed # 只重跑失败的
"""
import argparse
import json
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm
from common.path_utils import get_file_path
from _002__extract_information.__000__extract_graph_data_utils import (
build_alpaca_record,
extract_tcm_knowledge,
)
INPUT_DIR = get_file_path("_001__clawler/中药")
OUTPUT_DIR = get_file_path("_002__extract_information/中药图谱")
FAILED_LOG = os.path.join(OUTPUT_DIR, "_failed.log")
# 微调数据集用 jsonl 追加：一行一条，多线程各写一行不会互相截断，
# 中途崩了也只会丢最后一行，不会把前面写好的记录带坏。
ALPACA_JSONL = get_file_path("_002__extract_information/alpaca_herb_dataset.jsonl")
# 失败清单是多线程追加写的，加锁避免两行日志挤在一起
_failed_lock = threading.Lock()
# 数据集同理，一次只让一个线程追加
_alpaca_lock = threading.Lock()
def read_failed_names():
"""读失败清单里的中药名，用于 --retry-failed。"""
if not os.path.exists(FAILED_LOG):
return set()
names = set()
with open(FAILED_LOG, encoding="utf-8") as f:
for line in f:
line = line.strip()
if line:
names.add(line.split("\t")[0])
return names
def record_failure(name, reason):
"""把失败原因追加到清单。写成一行，方便之后 --retry-failed 挑出来重跑。"""
reason = " ".join(str(reason).split())[:200]
with _failed_lock:
with open(FAILED_LOG, "a", encoding="utf-8") as f:
f.write(f"{name}\t{reason}\n")
def append_alpaca(text, entities, relations):
"""把这一条抽取结果追加成一行 Alpaca 样本。
只追加、不去重 —— 断点续抽保证同一味药一轮里只会抽一次。真要补历史或
去重，跑 02_build_alpaca_dataset.py 从图谱 json 重建。
"""
record = build_alpaca_record(text, entities, relations)
line = json.dumps(record, ensure_ascii=False)
with _alpaca_lock:
with open(ALPACA_JSONL, "a", encoding="utf-8") as f:
f.write(line + "\n")
def extract_one(name, retries=2):
"""抽一味药并落盘。返回 (name, ok, note, 实体数, 关系数)。由线程池并发调用。
LLM 偶发超时/限流/返回非法 json，失败退避重试；仍然失败就记进清单，
不写 json —— 这样下次跑还会被当成待抽项，断点续抽才是对的。
"""
input_path = os.path.join(INPUT_DIR, f"{name}.txt")
output_path = os.path.join(OUTPUT_DIR, f"{name}.json")
try:
with open(input_path, encoding="utf-8") as f:
text = f.read()
except OSError as exc:
record_failure(name, f"读取失败: {exc}")
return name, False, "读取失败", 0, 0
if not text.strip():
record_failure(name, "源文件为空")
return name, False, "源文件为空", 0, 0
last_exc = None
for attempt in range(1, retries + 1):
try:
result = extract_tcm_knowledge(text)
break
except Exception as exc: # 网络、限流、json 解析失败都在这
last_exc = exc
if attempt < retries:
time.sleep(5 * attempt)
else:
record_failure(name, f"抽取失败: {last_exc}")
return name, False, "抽取失败", 0, 0
# parser 是 JsonOutputParser，正常返回 dict；模型偶尔会吐出 list 或别的形状
if not isinstance(result, dict):
record_failure(name, f"返回结构异常: {type(result).__name__}")
return name, False, "结构异常", 0, 0
entities = result.get("entities") or []
relations = result.get("relations") or []
# 单实体、无关系的文本，prompt 要求返回空结构，属于正常结果，照样落盘，
# 否则每次重跑都会重复调用这些药材。
payload = {
"name": name,
"source_file": os.path.relpath(input_path, get_file_path("")),
"entities": entities,
"relations": relations,
}
# 先写临时文件再改名，中断时不会留下半个 json 被误判成已完成
temp_path = f"{output_path}.{threading.get_ident()}.part"
with open(temp_path, "w", encoding="utf-8") as f:
json.dump(payload, f, ensure_ascii=False, indent=2)
os.replace(temp_path, output_path)
# 图谱 json 落盘之后再追加数据集：反过来的话，万一追加完崩了、json 还没写成，
# 续抽会把这味药重抽一遍，数据集里就多出一条重复样本。
append_alpaca(text, entities, relations)
note = f"{len(entities)}实体/{len(relations)}关系"
return name, True, note, len(entities), len(relations)
def build_todo(limit=0, retry_failed=False):
"""列出本次要抽的中药名。断点续抽的核心：已有 json 的直接排除。"""
names = sorted(
os.path.splitext(f)[0]
for f in os.listdir(INPUT_DIR)
if f.endswith(".txt")
)
done = {
os.path.splitext(f)[0]
for f in os.listdir(OUTPUT_DIR)
if f.endswith(".json")
}
todo = [n for n in names if n not in done]
if retry_failed:
# 失败的本来就没写 json，已经在 todo 里；这里是缩小范围只跑它们
failed = read_failed_names()
todo = [n for n in todo if n in failed]
if limit:
todo = todo[:limit]
return names, done, todo
def run(limit=0, workers=10, retry_failed=False):
os.makedirs(OUTPUT_DIR, exist_ok=True)
names, done, todo = build_todo(limit=limit, retry_failed=retry_failed)
print(f"源目录共 {len(names)} 味中药，已抽取 {len(done)} 味，本次待抽 {len(todo)} 味")
print(f"并发线程数: {workers}")
if not todo:
print("没有需要抽取的内容。")
return
ok_count = fail_count = 0
entity_total = relation_total = 0
with ThreadPoolExecutor(max_workers=workers) as pool:
futures = {pool.submit(extract_one, name): name for name in todo}
# tqdm 显示进度；每条结果用 write() 打，避免把进度条撕开
with tqdm(total=len(todo), desc="抽取中", unit="味", ncols=90) as bar:
for future in as_completed(futures):
name = futures[future]
try:
_, ok, note, n_entity, n_relation = future.result()
except Exception as exc: # 单条异常不该拖垮整轮
record_failure(name, f"未捕获异常: {exc}")
ok, note, n_entity, n_relation = False, "异常", 0, 0
if ok:
ok_count += 1
entity_total += n_entity
relation_total += n_relation
else:
fail_count += 1
bar.write(f" {'✓' if ok else '✗'} {name} {note}")
bar.update(1)
bar.set_postfix(成功=ok_count, 失败=fail_count)
print(f"\n完成 {ok_count} 味，失败 {fail_count} 味")
print(f"共抽出 {entity_total} 个实体、{relation_total} 条关系")
print(f"输出目录: {OUTPUT_DIR}")
print(f"微调数据集（追加 {ok_count} 条）: {ALPACA_JSONL}")
if fail_count:
print(f"失败清单: {FAILED_LOG}（可用 --retry-failed 只重跑这些）")
def main():
parser = argparse.ArgumentParser(description="抽取中药详情中的实体和关系")
parser.add_argument(
"--limit",
type=int,
default=0,
help="本次最多抽多少味，0 表示不限制（默认 0）",
)
parser.add_argument(
"--workers",
type=int,
default=40,
help="并发线程数（默认 10，设 1 就是单线程）",
)
parser.add_argument(
"--retry-failed",
action="store_true",
help="只重跑 _failed.log 里记过的中药",
)
args = parser.parse_args()
if args.workers < 1:
parser.error("--workers 至少为 1")
run(limit=args.limit, workers=args.workers, retry_failed=args.retry_failed)
if __name__ == "__main__":
main()