pipeline.py 17.3 KB
Newer Older
赵小蒙's avatar
赵小蒙 committed
1
2
3
4
5
# coding=utf8
import sys
import time
from urllib.parse import quote

许瑞's avatar
许瑞 committed
6
7
8
9
10
11
12
from magic_pdf.libs.commons import (
    read_file,
    join_path,
    parse_bucket_key,
    formatted_time,
    s3_image_save_path,
)
赵小蒙's avatar
赵小蒙 committed
13
14
from magic_pdf.libs.drop_reason import DropReason
from magic_pdf.libs.json_compressor import JsonCompressor
15
from magic_pdf.dict2md.mkcontent import mk_universal_format
16
from magic_pdf.pdf_parse_by_txt import parse_pdf_by_txt
赵小蒙's avatar
赵小蒙 committed
17
18
19
from magic_pdf.filter.pdf_classify_by_type import classify
from magic_pdf.filter.pdf_meta_scan import pdf_meta_scan
from loguru import logger
赵小蒙's avatar
赵小蒙 committed
20

许瑞's avatar
许瑞 committed
21
from magic_pdf.pdf_parse_for_train import parse_pdf_for_train
赵小蒙's avatar
赵小蒙 committed
22
from magic_pdf.spark.base import exception_handler, get_data_source
许瑞's avatar
许瑞 committed
23
from magic_pdf.train_utils.convert_to_train_format import convert_to_train_format
赵小蒙's avatar
赵小蒙 committed
24
from magic_pdf.spark.s3 import get_s3_config, get_s3_client
赵小蒙's avatar
赵小蒙 committed
25
26
27
28



def meta_scan(jso: dict, doc_layout_check=True) -> dict:
许瑞's avatar
许瑞 committed
29
    s3_pdf_path = jso.get("file_location")
赵小蒙's avatar
赵小蒙 committed
30
31
    s3_config = get_s3_config(s3_pdf_path)
    if doc_layout_check:
许瑞's avatar
许瑞 committed
32
33
34
35
36
        if (
            "doc_layout_result" not in jso
        ):  # 检测json中是存在模型数据,如果没有则需要跳过该pdf
            jso["need_drop"] = True
            jso["drop_reason"] = DropReason.MISS_DOC_LAYOUT_RESULT
赵小蒙's avatar
赵小蒙 committed
37
38
            return jso
    try:
39
        data_source = get_data_source(jso)
许瑞's avatar
许瑞 committed
40
        file_id = jso.get("file_id")
赵小蒙's avatar
赵小蒙 committed
41
        book_name = f"{data_source}/{file_id}"
赵小蒙's avatar
赵小蒙 committed
42
43
44
45
46
47
48
49
50

        # 首页存在超量drawing问题
        # special_pdf_list = ['zlib/zlib_21822650']
        # if book_name in special_pdf_list:
        #     jso['need_drop'] = True
        #     jso['drop_reason'] = DropReason.SPECIAL_PDF
        #     return jso

        start_time = time.time()  # 记录开始时间
许瑞's avatar
许瑞 committed
51
52
53
54
        logger.info(
            f"book_name is:{book_name},start_time is:{formatted_time(start_time)}",
            file=sys.stderr,
        )
赵小蒙's avatar
赵小蒙 committed
55
56
57
58
59
        file_content = read_file(s3_pdf_path, s3_config)
        read_file_time = int(time.time() - start_time)  # 计算执行时间

        start_time = time.time()  # 记录开始时间
        res = pdf_meta_scan(s3_pdf_path, file_content)
许瑞's avatar
许瑞 committed
60
61
62
63
64
        if res.get(
            "need_drop", False
        ):  # 如果返回的字典里有need_drop,则提取drop_reason并跳过本次解析
            jso["need_drop"] = True
            jso["drop_reason"] = res["drop_reason"]
赵小蒙's avatar
赵小蒙 committed
65
        else:  # 正常返回
许瑞's avatar
许瑞 committed
66
67
68
69
            jso["pdf_meta"] = res
            jso["content"] = ""
            jso["remark"] = ""
            jso["data_url"] = ""
赵小蒙's avatar
赵小蒙 committed
70
71
        end_time = time.time()  # 记录结束时间
        meta_scan_time = int(end_time - start_time)  # 计算执行时间
许瑞's avatar
许瑞 committed
72
73
74
75
76
77
        logger.info(
            f"book_name is:{book_name},end_time is:{formatted_time(end_time)},read_file_time is:{read_file_time},meta_scan_time is:{meta_scan_time}",
            file=sys.stderr,
        )
        jso["read_file_time"] = read_file_time
        jso["meta_scan_time"] = meta_scan_time
赵小蒙's avatar
赵小蒙 committed
78
79
80
81
82
83
    except Exception as e:
        jso = exception_handler(jso, e)
    return jso


def classify_by_type(jso: dict, debug_mode=False) -> dict:
许瑞's avatar
许瑞 committed
84
    # 检测debug开关
赵小蒙's avatar
赵小蒙 committed
85
86
    if debug_mode:
        pass
许瑞's avatar
许瑞 committed
87
88
    else:  # 如果debug没开,则检测是否有needdrop字段
        if jso.get("need_drop", False):
赵小蒙's avatar
赵小蒙 committed
89
90
91
            return jso
    # 开始正式逻辑
    try:
许瑞's avatar
许瑞 committed
92
        pdf_meta = jso.get("pdf_meta")
93
        data_source = get_data_source(jso)
许瑞's avatar
许瑞 committed
94
        file_id = jso.get("file_id")
赵小蒙's avatar
赵小蒙 committed
95
        book_name = f"{data_source}/{file_id}"
赵小蒙's avatar
赵小蒙 committed
96
97
98
99
        total_page = pdf_meta["total_page"]
        page_width = pdf_meta["page_width_pts"]
        page_height = pdf_meta["page_height_pts"]
        img_sz_list = pdf_meta["image_info_per_page"]
许瑞's avatar
许瑞 committed
100
101
102
103
        img_num_list = pdf_meta["imgs_per_page"]
        text_len_list = pdf_meta["text_len_per_page"]
        text_layout_list = pdf_meta["text_layout_per_page"]
        text_language = pdf_meta["text_language"]
赵小蒙's avatar
赵小蒙 committed
104
105
106
107
108
109
        # allow_language = ['zh', 'en']  # 允许的语言,目前只允许简中和英文的

        # if text_language not in allow_language:  # 如果语言不在允许的语言中,则drop
        #     jso['need_drop'] = True
        #     jso['drop_reason'] = DropReason.NOT_ALLOW_LANGUAGE
        #     return jso
许瑞's avatar
许瑞 committed
110
111
112
113
114
115
116
117
        pdf_path = pdf_meta["pdf_path"]
        is_encrypted = pdf_meta["is_encrypted"]
        is_needs_password = pdf_meta["is_needs_password"]
        if (
            is_encrypted or is_needs_password
        ):  # 加密的,需要密码的,没有页面的,都不处理
            jso["need_drop"] = True
            jso["drop_reason"] = DropReason.ENCRYPTED
赵小蒙's avatar
赵小蒙 committed
118
119
        else:
            start_time = time.time()  # 记录开始时间
许瑞's avatar
许瑞 committed
120
121
122
123
124
125
126
127
128
129
            is_text_pdf, results = classify(
                pdf_path,
                total_page,
                page_width,
                page_height,
                img_sz_list,
                text_len_list,
                img_num_list,
                text_layout_list,
            )
赵小蒙's avatar
赵小蒙 committed
130
131
            classify_time = int(time.time() - start_time)  # 计算执行时间
            if is_text_pdf:
许瑞's avatar
许瑞 committed
132
                pdf_meta["is_text_pdf"] = is_text_pdf
133
                jso["_pdf_type"] = "TXT"
许瑞's avatar
许瑞 committed
134
135
                jso["pdf_meta"] = pdf_meta
                jso["classify_time"] = classify_time
赵小蒙's avatar
赵小蒙 committed
136
137
                # print(json.dumps(pdf_meta, ensure_ascii=False))

许瑞's avatar
许瑞 committed
138
139
140
141
142
143
                allow_language = ["zh", "en"]  # 允许的语言,目前只允许简中和英文的
                if (
                    text_language not in allow_language
                ):  # 如果语言不在允许的语言中,则drop
                    jso["need_drop"] = True
                    jso["drop_reason"] = DropReason.NOT_ALLOW_LANGUAGE
赵小蒙's avatar
赵小蒙 committed
144
145
146
                    return jso
            else:
                # 先不drop
许瑞's avatar
许瑞 committed
147
                pdf_meta["is_text_pdf"] = is_text_pdf
148
                jso["_pdf_type"] = "OCR"
许瑞's avatar
许瑞 committed
149
150
                jso["pdf_meta"] = pdf_meta
                jso["classify_time"] = classify_time
151
152
                # jso["need_drop"] = True
                # jso["drop_reason"] = DropReason.NOT_IS_TEXT_PDF
赵小蒙's avatar
赵小蒙 committed
153
154
155
156
                extra_info = {"classify_rules": []}
                for condition, result in results.items():
                    if not result:
                        extra_info["classify_rules"].append(condition)
许瑞's avatar
许瑞 committed
157
                jso["extra_info"] = extra_info
赵小蒙's avatar
赵小蒙 committed
158
159
160
161
162
163
164
165
166
167

    except Exception as e:
        jso = exception_handler(jso, e)
    return jso


def save_tables_to_s3(jso: dict, debug_mode=False) -> dict:

    if debug_mode:
        pass
许瑞's avatar
许瑞 committed
168
169
170
171
172
173
    else:  # 如果debug没开,则检测是否有needdrop字段
        if jso.get("need_drop", False):
            logger.info(
                f"book_name is:{get_data_source(jso)}/{jso['file_id']} need drop",
                file=sys.stderr,
            )
赵小蒙's avatar
赵小蒙 committed
174
175
176
            jso["dropped"] = True
            return jso
    try:
177
        data_source = get_data_source(jso)
许瑞's avatar
许瑞 committed
178
        file_id = jso.get("file_id")
赵小蒙's avatar
赵小蒙 committed
179
        book_name = f"{data_source}/{file_id}"
许瑞's avatar
许瑞 committed
180
181
182
        title = jso.get("title")
        url_encode_title = quote(title, safe="")
        if data_source != "scihub":
赵小蒙's avatar
赵小蒙 committed
183
            return jso
许瑞's avatar
许瑞 committed
184
        pdf_intermediate_dict = jso["pdf_intermediate_dict"]
赵小蒙's avatar
赵小蒙 committed
185
186
187
188
        # 将 pdf_intermediate_dict 解压
        pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
        i = 0
        for page in pdf_intermediate_dict.values():
许瑞's avatar
许瑞 committed
189
190
            if page.get("tables"):
                if len(page["tables"]) > 0:
赵小蒙's avatar
赵小蒙 committed
191
                    j = 0
许瑞's avatar
许瑞 committed
192
                    for table in page["tables"]:
赵小蒙's avatar
赵小蒙 committed
193
                        if debug_mode:
许瑞's avatar
许瑞 committed
194
195
196
197
198
                            image_path = join_path(
                                "s3://mllm-raw-media/pdf2md_img/",
                                book_name,
                                table["image_path"],
                            )
赵小蒙's avatar
赵小蒙 committed
199
                        else:
许瑞's avatar
许瑞 committed
200
201
202
                            image_path = join_path(
                                "s3://mllm-raw-media/pdf2md_img/", table["image_path"]
                            )
赵小蒙's avatar
赵小蒙 committed
203

许瑞's avatar
许瑞 committed
204
                        if image_path.endswith(".jpg"):
赵小蒙's avatar
赵小蒙 committed
205
206
207
208
                            j += 1
                            s3_client = get_s3_client(image_path)
                            bucket_name, bucket_key = parse_bucket_key(image_path)
                            # 通过s3_client获取图片到内存
许瑞's avatar
许瑞 committed
209
210
211
                            image_bytes = s3_client.get_object(
                                Bucket=bucket_name, Key=bucket_key
                            )["Body"].read()
赵小蒙's avatar
赵小蒙 committed
212
213
                            # 保存图片到新的位置
                            if debug_mode:
许瑞's avatar
许瑞 committed
214
215
216
217
218
219
                                new_image_path = join_path(
                                    "s3://mllm-raw-media/pdf2md_img/table_new/",
                                    url_encode_title
                                    + "_"
                                    + table["image_path"].lstrip("tables/"),
                                )
赵小蒙's avatar
赵小蒙 committed
220
                            else:
许瑞's avatar
许瑞 committed
221
222
223
224
                                new_image_path = join_path(
                                    "s3://mllm-raw-media/pdf2md_img/table_new/",
                                    url_encode_title + f"_page{i}_{j}.jpg",
                                )
赵小蒙's avatar
赵小蒙 committed
225
226
227

                            logger.info(new_image_path, file=sys.stderr)
                            bucket_name, bucket_key = parse_bucket_key(new_image_path)
许瑞's avatar
许瑞 committed
228
229
230
                            s3_client.put_object(
                                Bucket=bucket_name, Key=bucket_key, Body=image_bytes
                            )
赵小蒙's avatar
赵小蒙 committed
231
232
233
234
235
236
237
238
239
240
241
242
243
244
                        else:
                            continue
            i += 1

        # 把无用的信息清空
        jso["doc_layout_result"] = ""
        jso["pdf_intermediate_dict"] = ""
        jso["pdf_meta"] = ""
    except Exception as e:
        jso = exception_handler(jso, e)
    return jso


def drop_needdrop_pdf(jso: dict) -> dict:
许瑞's avatar
许瑞 committed
245
246
247
248
249
    if jso.get("need_drop", False):
        logger.info(
            f"book_name is:{get_data_source(jso)}/{jso['file_id']} need drop",
            file=sys.stderr,
        )
赵小蒙's avatar
赵小蒙 committed
250
251
252
253
254
255
256
257
        jso["dropped"] = True
    return jso


def pdf_intermediate_dict_to_markdown(jso: dict, debug_mode=False) -> dict:

    if debug_mode:
        pass
许瑞's avatar
许瑞 committed
258
259
260
    else:  # 如果debug没开,则检测是否有needdrop字段
        if jso.get("need_drop", False):
            book_name = join_path(get_data_source(jso), jso["file_id"])
赵小蒙's avatar
赵小蒙 committed
261
262
263
264
            logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
            jso["dropped"] = True
            return jso
    try:
许瑞's avatar
许瑞 committed
265
        pdf_intermediate_dict = jso["pdf_intermediate_dict"]
赵小蒙's avatar
赵小蒙 committed
266
267
        # 将 pdf_intermediate_dict 解压
        pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
许瑞's avatar
许瑞 committed
268
269
270
        # markdown_content = mk_nlp_markdown(pdf_intermediate_dict)
        jso["content_list"] = mk_universal_format(pdf_intermediate_dict)
        # jso["content"] = markdown_content
271
        logger.info(f"book_name is:{get_data_source(jso)}/{jso['file_id']}")
赵小蒙's avatar
赵小蒙 committed
272
273
274
275
276
277
278
279
280
281
        # 把无用的信息清空
        jso["doc_layout_result"] = ""
        jso["pdf_intermediate_dict"] = ""
        jso["pdf_meta"] = ""
    except Exception as e:
        jso = exception_handler(jso, e)
    return jso


def parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
许瑞's avatar
许瑞 committed
282
    # 检测debug开关
赵小蒙's avatar
赵小蒙 committed
283
284
    if debug_mode:
        pass
许瑞's avatar
许瑞 committed
285
286
    else:  # 如果debug没开,则检测是否有needdrop字段
        if jso.get("need_drop", False):
赵小蒙's avatar
赵小蒙 committed
287
288
            return jso
    # 开始正式逻辑
许瑞's avatar
许瑞 committed
289
    s3_pdf_path = jso.get("file_location")
赵小蒙's avatar
赵小蒙 committed
290
    s3_config = get_s3_config(s3_pdf_path)
赵小蒙's avatar
赵小蒙 committed
291
    pdf_bytes = read_file(s3_pdf_path, s3_config)
许瑞's avatar
许瑞 committed
292
    model_output_json_list = jso.get("doc_layout_result")
293
    data_source = get_data_source(jso)
许瑞's avatar
许瑞 committed
294
    file_id = jso.get("file_id")
赵小蒙's avatar
赵小蒙 committed
295
    book_name = f"{data_source}/{file_id}"
赵小蒙's avatar
赵小蒙 committed
296

许瑞's avatar
许瑞 committed
297
    junk_img_bojids = jso["pdf_meta"]["junk_img_bojids"]
赵小蒙's avatar
赵小蒙 committed
298
299

    # 增加检测 max_svgs 数量的检测逻辑,如果 max_svgs 超过3000则drop
许瑞's avatar
许瑞 committed
300
    svgs_per_page_list = jso["pdf_meta"]["svgs_per_page"]
赵小蒙's avatar
赵小蒙 committed
301
302
    max_svgs = max(svgs_per_page_list)
    if max_svgs > 3000:
许瑞's avatar
许瑞 committed
303
304
        jso["need_drop"] = True
        jso["drop_reason"] = DropReason.HIGH_COMPUTATIONAL_lOAD_BY_SVGS
赵小蒙's avatar
赵小蒙 committed
305
306
    else:
        try:
赵小蒙's avatar
赵小蒙 committed
307
            save_path = s3_image_save_path
赵小蒙's avatar
赵小蒙 committed
308
309
310
            image_s3_config = get_s3_config(save_path)
            start_time = time.time()  # 记录开始时间
            # 先打印一下book_name和解析开始的时间
许瑞's avatar
许瑞 committed
311
312
313
314
            logger.info(
                f"book_name is:{book_name},start_time is:{formatted_time(start_time)}",
                file=sys.stderr,
            )
315
            pdf_info_dict = parse_pdf_by_txt(
赵小蒙's avatar
赵小蒙 committed
316
                pdf_bytes,
许瑞's avatar
许瑞 committed
317
318
319
320
321
322
323
324
325
326
327
328
329
330
                model_output_json_list,
                save_path,
                book_name,
                pdf_model_profile=None,
                image_s3_config=image_s3_config,
                start_page_id=start_page_id,
                junk_img_bojids=junk_img_bojids,
                debug_mode=debug_mode,
            )
            if pdf_info_dict.get(
                "need_drop", False
            ):  # 如果返回的字典里有need_drop,则提取drop_reason并跳过本次解析
                jso["need_drop"] = True
                jso["drop_reason"] = pdf_info_dict["drop_reason"]
赵小蒙's avatar
赵小蒙 committed
331
332
            else:  # 正常返回,将 pdf_info_dict 压缩并存储
                pdf_info_dict = JsonCompressor.compress_json(pdf_info_dict)
许瑞's avatar
许瑞 committed
333
                jso["pdf_intermediate_dict"] = pdf_info_dict
赵小蒙's avatar
赵小蒙 committed
334
335
336
            end_time = time.time()  # 记录完成时间
            parse_time = int(end_time - start_time)  # 计算执行时间
            # 解析完成后打印一下book_name和耗时
许瑞's avatar
许瑞 committed
337
338
339
340
341
            logger.info(
                f"book_name is:{book_name},end_time is:{formatted_time(end_time)},cost_time is:{parse_time}",
                file=sys.stderr,
            )
            jso["parse_time"] = parse_time
赵小蒙's avatar
赵小蒙 committed
342
343
344
345
        except Exception as e:
            jso = exception_handler(jso, e)
    return jso

许瑞's avatar
许瑞 committed
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410

def parse_pdf_for_model_train(jso: dict, start_page_id=0, debug_mode=False) -> dict:
    # 检测debug开关
    if debug_mode:
        pass
    else:  # 如果debug没开,则检测是否有needdrop字段
        if jso.get("need_drop", False):
            return jso
    # 开始正式逻辑
    s3_pdf_path = jso.get("file_location")
    s3_config = get_s3_config(s3_pdf_path)
    model_output_json_list = jso.get("doc_layout_result")
    data_source = get_data_source(jso)
    file_id = jso.get("file_id")
    book_name = f"{data_source}/{file_id}"

    # 1.23.22已修复
    # if debug_mode:
    #     pass
    # else:
    #     if book_name == "zlib/zlib_21929367":
    #         jso['need_drop'] = True
    #         jso['drop_reason'] = DropReason.SPECIAL_PDF
    #         return jso

    junk_img_bojids = jso["pdf_meta"]["junk_img_bojids"]
    # total_page = jso['pdf_meta']['total_page']

    # 增加检测 max_svgs 数量的检测逻辑,如果 max_svgs 超过3000则drop
    svgs_per_page_list = jso["pdf_meta"]["svgs_per_page"]
    max_svgs = max(svgs_per_page_list)
    if max_svgs > 3000:
        jso["need_drop"] = True
        jso["drop_reason"] = DropReason.HIGH_COMPUTATIONAL_lOAD_BY_SVGS
    # elif total_page > 1000:
    #     jso['need_drop'] = True
    #     jso['drop_reason'] = DropReason.HIGH_COMPUTATIONAL_lOAD_BY_TOTAL_PAGES
    else:
        try:
            save_path = s3_image_save_path
            image_s3_config = get_s3_config(save_path)
            start_time = time.time()  # 记录开始时间
            # 先打印一下book_name和解析开始的时间
            logger.info(
                f"book_name is:{book_name},start_time is:{formatted_time(start_time)}",
                file=sys.stderr,
            )
            pdf_info_dict = parse_pdf_for_train(
                s3_pdf_path,
                s3_config,
                model_output_json_list,
                save_path,
                book_name,
                pdf_model_profile=None,
                image_s3_config=image_s3_config,
                start_page_id=start_page_id,
                junk_img_bojids=junk_img_bojids,
                debug_mode=debug_mode,
            )
            if pdf_info_dict.get(
                "need_drop", False
            ):  # 如果返回的字典里有need_drop,则提取drop_reason并跳过本次解析
                jso["need_drop"] = True
                jso["drop_reason"] = pdf_info_dict["drop_reason"]
            else:  # 正常返回,将 pdf_info_dict 压缩并存储
许瑞's avatar
许瑞 committed
411
                jso["parsed_results"] = convert_to_train_format(pdf_info_dict)
许瑞's avatar
许瑞 committed
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
                pdf_info_dict = JsonCompressor.compress_json(pdf_info_dict)
                jso["pdf_intermediate_dict"] = pdf_info_dict
            end_time = time.time()  # 记录完成时间
            parse_time = int(end_time - start_time)  # 计算执行时间
            # 解析完成后打印一下book_name和耗时
            logger.info(
                f"book_name is:{book_name},end_time is:{formatted_time(end_time)},cost_time is:{parse_time}",
                file=sys.stderr,
            )
            jso["parse_time"] = parse_time
        except Exception as e:
            jso = exception_handler(jso, e)
    return jso


赵小蒙's avatar
赵小蒙 committed
427
428
429
430
431
432
"""
统一处理逻辑
1.先调用parse_pdf对文本类pdf进行处理
2.再调用ocr_dropped_parse_pdf,对之前drop的pdf进行处理
"""

433
434
435
436
# def uni_parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
#     jso = parse_pdf(jso, start_page_id=start_page_id, debug_mode=debug_mode)
#     jso = ocr_dropped_parse_pdf(jso, start_page_id=start_page_id, debug_mode=debug_mode)
#     return jso
赵小蒙's avatar
赵小蒙 committed
437
438


赵小蒙's avatar
赵小蒙 committed
439
440
if __name__ == "__main__":
    pass