magicpdf.py 9.62 KB
Newer Older
kernel.h@qq.com's avatar
kernel.h@qq.com committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
"""
这里实现2个click命令:
第一个:
 接收一个完整的s3路径,例如:s3://llm-pdf-text/pdf_ebook_and_paper/pre-clean-mm-markdown/v014/part-660420b490be-000008.jsonl?bytes=0,81350
    1)根据~/magic-pdf.json里的ak,sk等,构造s3cliReader读取到这个jsonl的对应行,返回json对象。
    2)根据Json对象里的pdf的s3路径获取到他的ak,sk,endpoint,构造出s3cliReader用来读取pdf
    3)从magic-pdf.json里读取到本地保存图片、Md等的临时目录位置,构造出LocalImageWriter,用来保存截图
    4)从magic-pdf.json里读取到本地保存图片、Md等的临时目录位置,构造出LocalIRdWriter,用来读写本地文件
    
    最后把以上步骤准备好的对象传入真正的解析API
    
第二个:
  接收1)pdf的本地路径。2)模型json文件(可选)。然后:
    1)根据~/magic-pdf.json读取到本地保存图片、md等临时目录的位置,构造出LocalImageWriter,用来保存截图
    2)从magic-pdf.json里读取到本地保存图片、Md等的临时目录位置,构造出LocalIRdWriter,用来读写本地文件
    3)根据约定,根据pdf本地路径,推导出pdf模型的json,并读入
    

效果:
python magicpdf.py --json  s3://llm-pdf-text/scihub/xxxx.json?bytes=0,81350 
python magicpdf.py --pdf  /home/llm/Downloads/xxxx.pdf --model /home/llm/Downloads/xxxx.json  或者 python magicpdf.py --pdf  /home/llm/Downloads/xxxx.pdf
"""

许瑞's avatar
许瑞 committed
24
25
import os
import json as json_parse
26
import sys
许瑞's avatar
许瑞 committed
27
import click
28
from loguru import logger
kernel.h@qq.com's avatar
kernel.h@qq.com committed
29
from pathlib import Path
赵小蒙's avatar
赵小蒙 committed
30
from magic_pdf.libs.version import __version__
31

32
from magic_pdf.libs.MakeContentConfig import DropMode
赵小蒙's avatar
赵小蒙 committed
33
from magic_pdf.libs.draw_bbox import draw_layout_bbox, draw_span_bbox
许瑞's avatar
许瑞 committed
34
from magic_pdf.pipe.UNIPipe import UNIPipe
许瑞's avatar
许瑞 committed
35
36
from magic_pdf.pipe.OCRPipe import OCRPipe
from magic_pdf.pipe.TXTPipe import TXTPipe
许瑞's avatar
许瑞 committed
37
38
39
40
41
42
43
from magic_pdf.libs.config_reader import get_s3_config
from magic_pdf.libs.path_utils import (
    parse_s3path,
    parse_s3_range_params,
    remove_non_official_s3_args,
)
from magic_pdf.libs.config_reader import get_local_dir
许瑞's avatar
许瑞 committed
44
from magic_pdf.rw.S3ReaderWriter import S3ReaderWriter
kernel.h@qq.com's avatar
kernel.h@qq.com committed
45
from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
许瑞's avatar
许瑞 committed
46
from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
kernel.h@qq.com's avatar
kernel.h@qq.com committed
47

48
49
50
parse_pdf_methods = click.Choice(["ocr", "txt", "auto"])


赵小蒙's avatar
赵小蒙 committed
51
def prepare_env(pdf_file_name, method):
许瑞's avatar
许瑞 committed
52
    local_parent_dir = os.path.join(
赵小蒙's avatar
赵小蒙 committed
53
        get_local_dir(), "magic-pdf", pdf_file_name, method
许瑞's avatar
许瑞 committed
54
55
    )

赵小蒙's avatar
赵小蒙 committed
56
    local_image_dir = os.path.join(str(local_parent_dir), "images")
57
    local_md_dir = local_parent_dir
许瑞's avatar
许瑞 committed
58
59
60
    os.makedirs(local_image_dir, exist_ok=True)
    os.makedirs(local_md_dir, exist_ok=True)
    return local_image_dir, local_md_dir
kernel.h@qq.com's avatar
kernel.h@qq.com committed
61
62


赵小蒙's avatar
赵小蒙 committed
63
def _do_parse(pdf_file_name, pdf_bytes, model_list, parse_method, image_writer, md_writer, image_dir, local_md_dir):
许瑞's avatar
许瑞 committed
64
    if parse_method == "auto":
赵小蒙's avatar
赵小蒙 committed
65
66
67
68
69
        jso_useful_key = {
            "_pdf_type": "",
            "model_list": model_list
        }
        pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer, is_debug=True)
许瑞's avatar
许瑞 committed
70
    elif parse_method == "txt":
71
        pipe = TXTPipe(pdf_bytes, model_list, image_writer, is_debug=True)
许瑞's avatar
许瑞 committed
72
    elif parse_method == "ocr":
73
        pipe = OCRPipe(pdf_bytes, model_list, image_writer, is_debug=True)
许瑞's avatar
许瑞 committed
74
75
    else:
        print("unknow parse method")
76
        sys.exit(1)
许瑞's avatar
许瑞 committed
77
78
79

    pipe.pipe_classify()
    pipe.pipe_parse()
赵小蒙's avatar
赵小蒙 committed
80
81
    pdf_info = pipe.pdf_mid_data['pdf_info']
    draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir)
赵小蒙's avatar
赵小蒙 committed
82
    draw_span_bbox(pdf_info, pdf_bytes, local_md_dir)
83

84
    md_content = pipe.pipe_mk_markdown(image_dir, drop_mode=DropMode.NONE)
许瑞's avatar
许瑞 committed
85
    md_writer.write(
kernel.h@qq.com's avatar
kernel.h@qq.com committed
86
        content=md_content, path=f"{pdf_file_name}.md", mode=AbsReaderWriter.MODE_TXT
许瑞's avatar
许瑞 committed
87
88
89
    )
    md_writer.write(
        content=json_parse.dumps(pipe.pdf_mid_data, ensure_ascii=False, indent=4),
kernel.h@qq.com's avatar
kernel.h@qq.com committed
90
        path=f"{pdf_file_name}.json",
许瑞's avatar
许瑞 committed
91
92
        mode=AbsReaderWriter.MODE_TXT,
    )
93
94

    content_list = pipe.pipe_mk_uni_format(image_dir, drop_mode=DropMode.NONE)
赵小蒙's avatar
赵小蒙 committed
95
96
97
    md_writer.write(
        str(content_list), f"{pdf_file_name}.txt", AbsReaderWriter.MODE_TXT
    )
许瑞's avatar
许瑞 committed
98
99


kernel.h@qq.com's avatar
kernel.h@qq.com committed
100
101
102
103
@click.group()
def cli():
    pass

许瑞's avatar
许瑞 committed
104

kernel.h@qq.com's avatar
kernel.h@qq.com committed
105
@cli.command()
赵小蒙's avatar
赵小蒙 committed
106
@click.version_option(__version__, "--version", "-v", help="显示版本信息")
许瑞's avatar
许瑞 committed
107
@click.option("--json", type=str, help="输入一个S3路径")
108
109
110
111
112
113
114
@click.option(
    "--method",
    type=parse_pdf_methods,
    help="指定解析方法。txt: 文本型 pdf 解析方法, ocr: 光学识别解析 pdf, auto: 程序智能选择解析方法",
    default="auto",
)
def json_command(json, method):
许瑞's avatar
许瑞 committed
115
116
    if not json.startswith("s3://"):
        print("usage: python magipdf.py --json s3://some_bucket/some_path")
117
        sys.exit(1)
许瑞's avatar
许瑞 committed
118
119
120
121
122
123
124
125

    def read_s3_path(s3path):
        bucket, key = parse_s3path(s3path)

        s3_ak, s3_sk, s3_endpoint = get_s3_config(bucket)
        s3_rw = S3ReaderWriter(
            s3_ak, s3_sk, s3_endpoint, "auto", remove_non_official_s3_args(s3path)
        )
126
        may_range_params = parse_s3_range_params(s3path)
许瑞's avatar
许瑞 committed
127
128
129
130
        if may_range_params is None or 2 != len(may_range_params):
            byte_start, byte_end = 0, None
        else:
            byte_start, byte_end = int(may_range_params[0]), int(may_range_params[1])
许瑞's avatar
许瑞 committed
131
            byte_end += byte_start - 1
许瑞's avatar
许瑞 committed
132
        return s3_rw.read_jsonl(
许瑞's avatar
许瑞 committed
133
134
135
136
            remove_non_official_s3_args(s3path),
            byte_start,
            byte_end,
            AbsReaderWriter.MODE_BIN,
许瑞's avatar
许瑞 committed
137
138
139
        )

    jso = json_parse.loads(read_s3_path(json).decode("utf-8"))
赵小蒙's avatar
赵小蒙 committed
140
141
142
    s3_file_path = jso.get("file_location")
    if s3_file_path is None:
        s3_file_path = jso.get("path")
kernel.h@qq.com's avatar
kernel.h@qq.com committed
143
144
    pdf_file_name = Path(s3_file_path).stem
    pdf_data = read_s3_path(s3_file_path)
赵小蒙's avatar
赵小蒙 committed
145
    local_image_dir, local_md_dir = prepare_env(pdf_file_name, method)
kernel.h@qq.com's avatar
kernel.h@qq.com committed
146
    
许瑞's avatar
许瑞 committed
147
148
149
    local_image_rw, local_md_rw = DiskReaderWriter(local_image_dir), DiskReaderWriter(
        local_md_dir
    )
许瑞's avatar
许瑞 committed
150

许瑞's avatar
许瑞 committed
151
    _do_parse(
kernel.h@qq.com's avatar
kernel.h@qq.com committed
152
        pdf_file_name,
许瑞's avatar
许瑞 committed
153
        pdf_data,
许瑞's avatar
许瑞 committed
154
        jso["doc_layout_result"],
许瑞's avatar
许瑞 committed
155
156
157
        method,
        local_image_rw,
        local_md_rw,
158
        os.path.basename(local_image_dir),
赵小蒙's avatar
赵小蒙 committed
159
        local_md_dir
许瑞's avatar
许瑞 committed
160
    )
许瑞's avatar
许瑞 committed
161

kernel.h@qq.com's avatar
kernel.h@qq.com committed
162

赵小蒙's avatar
赵小蒙 committed
163
164
165
166
167
168
169
170
171
172
173
174
@cli.command()
@click.version_option(__version__, "--version", "-v", help="显示版本信息")
@click.option("--local_json", type=str, help="输入一个本地jsonl路径")
@click.option(
    "--method",
    type=parse_pdf_methods,
    help="指定解析方法。txt: 文本型 pdf 解析方法, ocr: 光学识别解析 pdf, auto: 程序智能选择解析方法",
    default="auto",
)
def local_json_command(local_json, method):
    def read_s3_path(s3path):
        bucket, key = parse_s3path(s3path)
赵小蒙's avatar
赵小蒙 committed
175

赵小蒙's avatar
赵小蒙 committed
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
        s3_ak, s3_sk, s3_endpoint = get_s3_config(bucket)
        s3_rw = S3ReaderWriter(
            s3_ak, s3_sk, s3_endpoint, "auto", remove_non_official_s3_args(s3path)
        )
        may_range_params = parse_s3_range_params(s3path)
        if may_range_params is None or 2 != len(may_range_params):
            byte_start, byte_end = 0, None
        else:
            byte_start, byte_end = int(may_range_params[0]), int(may_range_params[1])
            byte_end += byte_start - 1
        return s3_rw.read_jsonl(
            remove_non_official_s3_args(s3path),
            byte_start,
            byte_end,
            AbsReaderWriter.MODE_BIN,
        )

    with open(local_json, "r", encoding="utf-8") as f:
        for json_line in f:
            jso = json_parse.loads(json_line)

            s3_file_path = jso.get("file_location")
            if s3_file_path is None:
                s3_file_path = jso.get("path")
            pdf_file_name = Path(s3_file_path).stem
            pdf_data = read_s3_path(s3_file_path)
            local_image_dir, local_md_dir = prepare_env(pdf_file_name, method)

            local_image_rw, local_md_rw = DiskReaderWriter(local_image_dir), DiskReaderWriter(
                local_md_dir
赵小蒙's avatar
赵小蒙 committed
206
207
            )

赵小蒙's avatar
赵小蒙 committed
208
209
210
211
212
213
214
215
216
217
            _do_parse(
                pdf_file_name,
                pdf_data,
                jso["doc_layout_result"],
                method,
                local_image_rw,
                local_md_rw,
                os.path.basename(local_image_dir),
                local_md_dir
            )
赵小蒙's avatar
赵小蒙 committed
218
219


kernel.h@qq.com's avatar
kernel.h@qq.com committed
220
@cli.command()
赵小蒙's avatar
赵小蒙 committed
221
@click.version_option(__version__, "--version", "-v", help="显示版本信息")
许瑞's avatar
许瑞 committed
222
223
224
225
@click.option(
    "--pdf", type=click.Path(exists=True), required=True, help="PDF文件的路径"
)
@click.option("--model", type=click.Path(exists=True), help="模型的路径")
226
227
228
229
230
231
232
@click.option(
    "--method",
    type=parse_pdf_methods,
    help="指定解析方法。txt: 文本型 pdf 解析方法, ocr: 光学识别解析 pdf, auto: 程序智能选择解析方法",
    default="auto",
)
def pdf_command(pdf, model, method):
kernel.h@qq.com's avatar
kernel.h@qq.com committed
233
    # 这里处理pdf和模型相关的逻辑
许瑞's avatar
许瑞 committed
234
235
236
    if model is None:
        model = pdf.replace(".pdf", ".json")
        if not os.path.exists(model):
237
238
            print(f"make sure json {model} existed and place under {os.path.dirname(pdf)}", file=sys.stderr)
            exit(1)
许瑞's avatar
许瑞 committed
239
240
241

    def read_fn(path):
        disk_rw = DiskReaderWriter(os.path.dirname(path))
许瑞's avatar
许瑞 committed
242
        return disk_rw.read(os.path.basename(path), AbsReaderWriter.MODE_BIN)
许瑞's avatar
许瑞 committed
243
244
245

    pdf_data = read_fn(pdf)
    jso = json_parse.loads(read_fn(model).decode("utf-8"))
kernel.h@qq.com's avatar
kernel.h@qq.com committed
246
    pdf_file_name = Path(pdf).stem
赵小蒙's avatar
赵小蒙 committed
247
    local_image_dir, local_md_dir = prepare_env(pdf_file_name, method)
许瑞's avatar
许瑞 committed
248
249
250
251
    local_image_rw, local_md_rw = DiskReaderWriter(local_image_dir), DiskReaderWriter(
        local_md_dir
    )
    _do_parse(
kernel.h@qq.com's avatar
kernel.h@qq.com committed
252
        pdf_file_name,
许瑞's avatar
许瑞 committed
253
        pdf_data,
kernel.h@qq.com's avatar
kernel.h@qq.com committed
254
        jso,
许瑞's avatar
许瑞 committed
255
256
257
        method,
        local_image_rw,
        local_md_rw,
258
        os.path.basename(local_image_dir),
赵小蒙's avatar
赵小蒙 committed
259
        local_md_dir
许瑞's avatar
许瑞 committed
260
    )
许瑞's avatar
许瑞 committed
261

kernel.h@qq.com's avatar
kernel.h@qq.com committed
262

赵小蒙's avatar
赵小蒙 committed
263

许瑞's avatar
许瑞 committed
264
265
if __name__ == "__main__":
    """
许瑞's avatar
许瑞 committed
266
    python magic_pdf/cli/magicpdf.py json-command --json s3://llm-pdf-text/pdf_ebook_and_paper/manual/v001/part-660407a28beb-000002.jsonl?bytes=0,63551
许瑞's avatar
许瑞 committed
267
    """
kernel.h@qq.com's avatar
kernel.h@qq.com committed
268
    cli()