1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
| #!/usr/bin/env python3
"""
PDF → MinerU 转 Markdown → 本地大模型按规则提取 JSON
"""
import argparse
import json
import re
import subprocess
import sys
from pathlib import Path
def query_llm_for_json(rule: str, text: str, model: str = "qwen2.5:14b") -> dict:
"""
用本地 Ollama 大模型根据规则从文本中提取信息，返回解析后的 JSON。
:param rule: 提示语/规则，例如「从以下文本中提取所有密评人员的姓名、是否有证书...」
:param text: 待分析的文本（如 Markdown 内容）
:param model: Ollama 模型名，默认 qwen2.5:14b
:return: 解析得到的 dict，若解析失败则抛出或返回空结构
"""
cmd = ["ollama", "run", model, rule]
try:
proc = subprocess.run(
cmd,
input=text.encode("utf-8"),
capture_output=True,
timeout=300,
cwd=None,
)
out = proc.stdout.decode("utf-8", errors="replace").strip()
err = proc.stderr.decode("utf-8", errors="replace").strip()
if err:
sys.stderr.write(f"ollama stderr: {err}\n")
except subprocess.TimeoutExpired:
raise TimeoutError(f"Ollama 调用超时 (model={model})")
except FileNotFoundError:
raise RuntimeError("未找到 ollama 命令，请先安装并确保在 PATH 中")
# 尽量从输出中剥离出纯 JSON
json_str = _extract_json_string(out)
if not json_str:
raise ValueError(f"无法从模型输出中解析 JSON。原始输出:\n{out[:500]}")
return json.loads(json_str)
def _extract_json_string(raw: str) -> str | None:
"""从模型输出中提取第一个完整的 JSON 字符串（支持 ```json ... ``` 包裹）。"""
raw = raw.strip()
# 尝试 ```json ... ``` 代码块
m = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", raw)
if m:
return m.group(1).strip()
# 尝试找 { ... } 包裹的 JSON
start = raw.find("{")
if start == -1:
return None
depth = 0
for i in range(start, len(raw)):
if raw[i] == "{":
depth += 1
elif raw[i] == "}":
depth -= 1
if depth == 0:
return raw[start : i + 1]
return None
def pdf_to_markdown_path(pdf_path: Path, output_dir: Path) -> Path:
"""
根据 MinerU 的产出规则，得到 PDF 对应的 Markdown 文件路径。
约定：output_dir / {pdf_stem} / hybrid_auto / {pdf_stem}.md
"""
stem = pdf_path.stem
return output_dir / stem / "hybrid_auto" / f"{stem}.md"
def run_mineru(pdf_path: Path, output_dir: Path) -> Path:
"""
调用 mineru 将 PDF 转为 Markdown，返回生成的 .md 文件路径。
"""
pdf_path = pdf_path.resolve()
output_dir = output_dir.resolve()
if not pdf_path.is_file():
raise FileNotFoundError(f"PDF 不存在: {pdf_path}")
cmd = ["mineru", "-p", str(pdf_path), "-o", str(output_dir)]
subprocess.run(cmd, check=True, timeout=600)
md_path = pdf_to_markdown_path(pdf_path, output_dir)
if not md_path.is_file():
raise FileNotFoundError(f"MinerU 未生成预期 Markdown: {md_path}")
return md_path
def extract_from_pdf(
pdf_path: str | Path,
rule: str,
output_dir: str | Path = ".",
model: str = "qwen2.5:14b",
) -> dict:
"""
一次调用：传入 PDF 路径和提取规则，先 MinerU 转 Markdown，再用大模型按规则提取 JSON。
:param pdf_path: PDF 文件路径
:param rule: 给大模型的提示语/规则，例如「从以下文本中提取所有密评人员的姓名、是否有证书、证书成绩、考试通过时间，提取后直接返回json数据给我，直接给我json字符串，不要有任何别的信息」
:param output_dir: MinerU 输出目录，默认当前目录
:param model: Ollama 模型名
:return: 大模型返回并解析后的 JSON（dict）
"""
pdf_path = Path(pdf_path)
output_dir = Path(output_dir)
md_path = run_mineru(pdf_path, output_dir)
text = md_path.read_text(encoding="utf-8", errors="replace")
return query_llm_for_json(rule, text, model=model)
# 预设规则示例，之后可在此扩展不同话术
DEFAULT_RULE = (
"从以下文本中提取所有密评人员的姓名、是否有证书、证书成绩、考试通过时间，"
"提取后直接返回json数据给我，直接给我json字符串，不要有任何别的信息"
)
def main():
parser = argparse.ArgumentParser(description="PDF → MinerU → 大模型提取 JSON")
parser.add_argument("pdf", type=Path, nargs="?", help="PDF 文件路径（使用 --only-query 时可不传）")
parser.add_argument("-o", "--output-dir", type=Path, default=Path("."), help="MinerU 输出目录，默认当前目录")
parser.add_argument(
"-r",
"--rule",
type=str,
default=DEFAULT_RULE,
help="大模型提取规则/话术，默认提取密评人员姓名、证书、成绩、通过时间",
)
parser.add_argument("-m", "--model", type=str, default="qwen2.5:14b", help="Ollama 模型名")
parser.add_argument("--only-query", action="store_true", help="仅对已有 Markdown 做一次大模型查询（需同时传 --md）")
parser.add_argument("--md", type=Path, help="已有 Markdown 文件路径（与 --only-query 一起使用）")
args = parser.parse_args()
if args.only_query:
if not args.md or not args.md.is_file():
parser.error("使用 --only-query 时必须指定已存在的 --md 文件")
text = args.md.read_text(encoding="utf-8", errors="replace")
result = query_llm_for_json(args.rule, text, model=args.model)
else:
if not args.pdf:
parser.error("未使用 --only-query 时必须指定 PDF 文件路径")
result = extract_from_pdf(
args.pdf,
args.rule,
output_dir=args.output_dir,
model=args.model,
)
print(json.dumps(result, ensure_ascii=False, indent=2))
return 0
if __name__ == "__main__":
sys.exit(main() or 0)
|
