magicpdf.py 7.17 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
许瑞's avatar
许瑞 committed
26
import click
27
from loguru import logger
kernel.h@qq.com's avatar
kernel.h@qq.com committed
28
from pathlib import Path
29

赵小蒙's avatar
赵小蒙 committed
30
from magic_pdf.libs.draw_bbox import draw_layout_bbox
许瑞's avatar
许瑞 committed
31
from magic_pdf.pipe.UNIPipe import UNIPipe
许瑞's avatar
许瑞 committed
32
33
from magic_pdf.pipe.OCRPipe import OCRPipe
from magic_pdf.pipe.TXTPipe import TXTPipe
许瑞's avatar
许瑞 committed
34
35
36
37
38
39
40
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
41
from magic_pdf.rw.S3ReaderWriter import S3ReaderWriter
kernel.h@qq.com's avatar
kernel.h@qq.com committed
42
from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
许瑞's avatar
许瑞 committed
43
from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
kernel.h@qq.com's avatar
kernel.h@qq.com committed
44

45
46
47
parse_pdf_methods = click.Choice(["ocr", "txt", "auto"])


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

    local_image_dir = os.path.join(local_parent_dir, "images")
54
    local_md_dir = local_parent_dir
许瑞's avatar
许瑞 committed
55
56
57
    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
58
59


赵小蒙's avatar
赵小蒙 committed
60
def _do_parse(pdf_file_name, pdf_bytes, model_list, parse_method, image_writer, md_writer, image_dir, local_md_dir):
许瑞's avatar
许瑞 committed
61
62
63
64
65
66
67
68
69
70
71
72
    if parse_method == "auto":
        pipe = UNIPipe(pdf_bytes, model_list, image_writer, image_dir, is_debug=True)
    elif parse_method == "txt":
        pipe = TXTPipe(pdf_bytes, model_list, image_writer, image_dir, is_debug=True)
    elif parse_method == "ocr":
        pipe = OCRPipe(pdf_bytes, model_list, image_writer, image_dir, is_debug=True)
    else:
        print("unknow parse method")
        os.exit(1)

    pipe.pipe_classify()
    pipe.pipe_parse()
赵小蒙's avatar
赵小蒙 committed
73
74
    pdf_info = pipe.pdf_mid_data['pdf_info']
    draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir)
许瑞's avatar
许瑞 committed
75
    md_content = pipe.pipe_mk_markdown()
kernel.h@qq.com's avatar
kernel.h@qq.com committed
76
    #part_file_name = datetime.now().strftime("%H-%M-%S")
许瑞's avatar
许瑞 committed
77
    md_writer.write(
kernel.h@qq.com's avatar
kernel.h@qq.com committed
78
        content=md_content, path=f"{pdf_file_name}.md", mode=AbsReaderWriter.MODE_TXT
许瑞's avatar
许瑞 committed
79
80
81
    )
    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
82
        path=f"{pdf_file_name}.json",
许瑞's avatar
许瑞 committed
83
84
        mode=AbsReaderWriter.MODE_TXT,
    )
85
86
87
88
89
90
91
    # try:
    #     content_list = pipe.pipe_mk_uni_format()
    # except Exception as e:
    #     logger.exception(e)
    # md_writer.write(
    #     str(content_list), f"{part_file_name}.txt", AbsReaderWriter.MODE_TXT
    # )
许瑞's avatar
许瑞 committed
92
93


kernel.h@qq.com's avatar
kernel.h@qq.com committed
94
95
96
97
@click.group()
def cli():
    pass

许瑞's avatar
许瑞 committed
98

kernel.h@qq.com's avatar
kernel.h@qq.com committed
99
@cli.command()
许瑞's avatar
许瑞 committed
100
@click.option("--json", type=str, help="输入一个S3路径")
101
102
103
104
105
106
107
@click.option(
    "--method",
    type=parse_pdf_methods,
    help="指定解析方法。txt: 文本型 pdf 解析方法, ocr: 光学识别解析 pdf, auto: 程序智能选择解析方法",
    default="auto",
)
def json_command(json, method):
许瑞's avatar
许瑞 committed
108
109
110
111
112
113
114
115
116
117
118
    if not json.startswith("s3://"):
        print("usage: python magipdf.py --json s3://some_bucket/some_path")
        os.exit(1)

    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)
        )
119
        may_range_params = parse_s3_range_params(s3path)
许瑞's avatar
许瑞 committed
120
121
122
123
        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
124
            byte_end += byte_start - 1
许瑞's avatar
许瑞 committed
125
        return s3_rw.read_jsonl(
许瑞's avatar
许瑞 committed
126
127
128
129
            remove_non_official_s3_args(s3path),
            byte_start,
            byte_end,
            AbsReaderWriter.MODE_BIN,
许瑞's avatar
许瑞 committed
130
131
132
        )

    jso = json_parse.loads(read_s3_path(json).decode("utf-8"))
kernel.h@qq.com's avatar
kernel.h@qq.com committed
133
134
135
    s3_file_path = jso["file_location"]
    pdf_file_name = Path(s3_file_path).stem
    pdf_data = read_s3_path(s3_file_path)
赵小蒙's avatar
赵小蒙 committed
136
    local_image_dir, local_md_dir = prepare_env(pdf_file_name, method)
kernel.h@qq.com's avatar
kernel.h@qq.com committed
137
    
许瑞's avatar
许瑞 committed
138
139
140
    local_image_rw, local_md_rw = DiskReaderWriter(local_image_dir), DiskReaderWriter(
        local_md_dir
    )
许瑞's avatar
许瑞 committed
141

许瑞's avatar
许瑞 committed
142
    _do_parse(
kernel.h@qq.com's avatar
kernel.h@qq.com committed
143
        pdf_file_name,
许瑞's avatar
许瑞 committed
144
        pdf_data,
许瑞's avatar
许瑞 committed
145
        jso["doc_layout_result"],
许瑞's avatar
许瑞 committed
146
147
148
        method,
        local_image_rw,
        local_md_rw,
149
        os.path.basename(local_image_dir),
赵小蒙's avatar
赵小蒙 committed
150
        local_md_dir
许瑞's avatar
许瑞 committed
151
    )
许瑞's avatar
许瑞 committed
152

kernel.h@qq.com's avatar
kernel.h@qq.com committed
153
154

@cli.command()
许瑞's avatar
许瑞 committed
155
156
157
158
@click.option(
    "--pdf", type=click.Path(exists=True), required=True, help="PDF文件的路径"
)
@click.option("--model", type=click.Path(exists=True), help="模型的路径")
159
160
161
162
163
164
165
@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
166
    # 这里处理pdf和模型相关的逻辑
许瑞's avatar
许瑞 committed
167
168
169
170
    if model is None:
        model = pdf.replace(".pdf", ".json")
        if not os.path.exists(model):
            print(f"make sure json file existed and place under {os.dirname(pdf)}")
赵小蒙's avatar
赵小蒙 committed
171
            os.exit(1)
许瑞's avatar
许瑞 committed
172
173
174

    def read_fn(path):
        disk_rw = DiskReaderWriter(os.path.dirname(path))
许瑞's avatar
许瑞 committed
175
        return disk_rw.read(os.path.basename(path), AbsReaderWriter.MODE_BIN)
许瑞's avatar
许瑞 committed
176
177
178

    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
179
    pdf_file_name = Path(pdf).stem
赵小蒙's avatar
赵小蒙 committed
180
    local_image_dir, local_md_dir = prepare_env(pdf_file_name, method)
许瑞's avatar
许瑞 committed
181
182
183
184
    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
185
        pdf_file_name,
许瑞's avatar
许瑞 committed
186
        pdf_data,
kernel.h@qq.com's avatar
kernel.h@qq.com committed
187
        jso,
许瑞's avatar
许瑞 committed
188
189
190
        method,
        local_image_rw,
        local_md_rw,
191
        os.path.basename(local_image_dir),
赵小蒙's avatar
赵小蒙 committed
192
        local_md_dir
许瑞's avatar
许瑞 committed
193
    )
许瑞's avatar
许瑞 committed
194

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

许瑞's avatar
许瑞 committed
196
197
if __name__ == "__main__":
    """
许瑞's avatar
许瑞 committed
198
    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
199
    """
kernel.h@qq.com's avatar
kernel.h@qq.com committed
200
    cli()