magic-pdf 0.5.4__py3-none-any.whl

magic_pdf/pre_proc/remove_rotate_bbox.py

@@ -0,0 +1,188 @@
+import math
+
+from magic_pdf.libs.boxbase import is_vbox_on_side
+from magic_pdf.libs.drop_tag import EMPTY_SIDE_BLOCK, ROTATE_TEXT, VERTICAL_TEXT
+
+
+def detect_non_horizontal_texts(result_dict):
+    """
+    This function detects watermarks and vertical margin notes in the document.
+
+    Watermarks are identified by finding blocks with the same coordinates and frequently occurring identical texts across multiple pages.
+    If these conditions are met, the blocks are highly likely to be watermarks, as opposed to headers or footers, which can change from page to page.
+    If the direction of these blocks is not horizontal, they are definitely considered to be watermarks.
+
+    Vertical margin notes are identified by finding blocks with the same coordinates and frequently occurring identical texts across multiple pages.
+    If these conditions are met, the blocks are highly likely to be vertical margin notes, which typically appear on the left and right sides of the page.
+    If the direction of these blocks is vertical, they are definitely considered to be vertical margin notes.
+
+
+    Parameters
+    ----------
+    result_dict : dict
+        The result dictionary.
+
+    Returns
+    -------
+    result_dict : dict
+        The updated result dictionary.
+    """
+    # Dictionary to store information about potential watermarks
+    potential_watermarks = {}
+    potential_margin_notes = {}
+
+    for page_id, page_content in result_dict.items():
+        if page_id.startswith("page_"):
+            for block_id, block_data in page_content.items():
+                if block_id.startswith("block_"):
+                    if "dir" in block_data:
+                        coordinates_text = (block_data["bbox"], block_data["text"])  # Tuple of coordinates and text
+
+                        angle = math.atan2(block_data["dir"][1], block_data["dir"][0])
+                        angle = abs(math.degrees(angle))
+
+                        if angle > 5 and angle < 85:  # Check if direction is watermarks
+                            if coordinates_text in potential_watermarks:
+                                potential_watermarks[coordinates_text] += 1
+                            else:
+                                potential_watermarks[coordinates_text] = 1
+
+                        if angle > 85 and angle < 105:  # Check if direction is vertical
+                            if coordinates_text in potential_margin_notes:
+                                potential_margin_notes[coordinates_text] += 1  # Increment count
+                            else:
+                                potential_margin_notes[coordinates_text] = 1  # Initialize count
+
+    # Identify watermarks by finding entries with counts higher than a threshold (e.g., appearing on more than half of the pages)
+    watermark_threshold = len(result_dict) // 2
+    watermarks = {k: v for k, v in potential_watermarks.items() if v > watermark_threshold}
+
+    # Identify margin notes by finding entries with counts higher than a threshold (e.g., appearing on more than half of the pages)
+    margin_note_threshold = len(result_dict) // 2
+    margin_notes = {k: v for k, v in potential_margin_notes.items() if v > margin_note_threshold}
+
+    # Add watermark information to the result dictionary
+    for page_id, blocks in result_dict.items():
+        if page_id.startswith("page_"):
+            for block_id, block_data in blocks.items():
+                coordinates_text = (block_data["bbox"], block_data["text"])
+                if coordinates_text in watermarks:
+                    block_data["is_watermark"] = 1
+                else:
+                    block_data["is_watermark"] = 0
+
+                if coordinates_text in margin_notes:
+                    block_data["is_vertical_margin_note"] = 1
+                else:
+                    block_data["is_vertical_margin_note"] = 0
+
+    return result_dict
+
+
+"""
+1. 当一个block里全部文字都不是dir=(1,0)，这个block整体去掉
+2. 当一个block里全部文字都是dir=(1,0)，但是每行只有一个字，这个block整体去掉。这个block必须出现在页面的四周，否则不去掉
+"""
+import re
+
+def __is_a_word(sentence):
+    # 如果输入是中文并且长度为1，则返回True
+    if re.fullmatch(r'[\u4e00-\u9fa5]', sentence):
+        return True
+    # 判断是否为单个英文单词或字符（包括ASCII标点）
+    elif re.fullmatch(r'[a-zA-Z0-9]+', sentence) and len(sentence) <=2:
+        return True
+    else:
+        return False
+
+
+def __get_text_color(num):
+    """获取字体的颜色RGB值"""
+    blue = num & 255
+    green = (num >> 8) & 255
+    red = (num >> 16) & 255
+    return red, green, blue
+
+
+def __is_empty_side_box(text_block):
+    """
+    是否是边缘上的空白没有任何内容的block
+    """
+    for line in text_block['lines']:
+        for span in line['spans']:
+            font_color = span['color']
+            r,g,b = __get_text_color(font_color)
+            if len(span['text'].strip())>0 and (r,g,b)!=(255,255,255):
+                return False
+            
+    return True
+
+
+def remove_rotate_side_textblock(pymu_text_block, page_width, page_height):
+    """
+    返回删除了垂直，水印，旋转的textblock
+    删除的内容打上tag返回
+    """
+    removed_text_block = []
+    
+    for i, block in enumerate(pymu_text_block): # 格式参考test/assets/papre/pymu_textblocks.json
+        lines = block['lines']
+        block_bbox = block['bbox']
+        if not is_vbox_on_side(block_bbox, page_width, page_height, 0.2): # 保证这些box必须在页面的两边
+           continue
+        
+        if all([__is_a_word(line['spans'][0]["text"]) for line in lines if len(line['spans'])>0]) and len(lines)>1 and all([len(line['spans'])==1 for line in lines]):
+            is_box_valign = (len(set([int(line['spans'][0]['bbox'][0] ) for line in lines if len(line['spans'])>0]))==1) and (len([int(line['spans'][0]['bbox'][0] ) for line in lines if len(line['spans'])>0])>1)  # 测试bbox在垂直方向是不是x0都相等，也就是在垂直方向排列.同时必须大于等于2个字
+            
+            if is_box_valign:
+                block['tag'] = VERTICAL_TEXT
+                removed_text_block.append(block)
+                continue
+        
+        for line in lines:
+            if line['dir']!=(1,0):
+                block['tag'] = ROTATE_TEXT
+                removed_text_block.append(block) # 只要有一个line不是dir=(1,0)，就把整个block都删掉
+                break
+        
+    for block in removed_text_block:
+        pymu_text_block.remove(block)
+    
+    return pymu_text_block, removed_text_block
+
+def get_side_boundry(rotate_bbox, page_width, page_height):
+    """
+    根据rotate_bbox，返回页面的左右正文边界
+    """
+    left_x = 0
+    right_x = page_width
+    for x in rotate_bbox:
+        box = x['bbox']
+        if box[2]<page_width/2:
+            left_x = max(left_x, box[2])
+        else:
+            right_x = min(right_x, box[0])
+            
+    return left_x+1, right_x-1
+
+
+def remove_side_blank_block(pymu_text_block, page_width, page_height):
+    """
+    删除页面两侧的空白block
+    """
+    removed_text_block = []
+    
+    for i, block in enumerate(pymu_text_block): # 格式参考test/assets/papre/pymu_textblocks.json
+        block_bbox = block['bbox']
+        if not is_vbox_on_side(block_bbox, page_width, page_height, 0.2): # 保证这些box必须在页面的两边
+           continue
+            
+        if __is_empty_side_box(block):
+            block['tag'] = EMPTY_SIDE_BLOCK
+            removed_text_block.append(block)
+            continue
+        
+    for block in removed_text_block:
+        pymu_text_block.remove(block)
+    
+    return pymu_text_block, removed_text_block

magic_pdf/pre_proc/resolve_bbox_conflict.py

@@ -0,0 +1,191 @@
+"""
+从pdf里提取出来api给出的bbox,然后根据重叠情况做出取舍
+1. 首先去掉出现在图片上的bbox，图片包括表格和图片
+2. 然后去掉出现在文字blcok上的图片bbox
+"""
+
+from magic_pdf.libs.boxbase import _is_in, _is_in_or_part_overlap, _is_left_overlap
+from magic_pdf.libs.drop_tag import ON_IMAGE_TEXT, ON_TABLE_TEXT
+
+
+def resolve_bbox_overlap_conflict(images: list, tables: list, interline_equations: list, inline_equations: list,
+                                  text_raw_blocks: list):
+    """
+    text_raw_blocks结构是从pymupdf里直接取到的结构，具体样例参考test/assets/papre/pymu_textblocks.json
+    当下采用一种粗暴的方式：
+    1. 去掉图片上的公式
+    2. 去掉table上的公式
+    2. 图片和文字block部分重叠，首先丢弃图片
+    3. 图片和图片重叠，修改图片的bbox，使得图片不重叠(暂时没这么做，先把图片都扔掉)
+    4. 去掉文字bbox里位于图片、表格上的文字（一定要完全在图、表内部）
+    5. 去掉表格上的文字
+    """
+    text_block_removed = []
+    images_backup = []
+
+    # 去掉位于图片上的文字block
+    for image_box in images:
+        for text_block in text_raw_blocks:
+            text_bbox = text_block["bbox"]
+            if _is_in(text_bbox, image_box):
+                text_block['tag'] = ON_IMAGE_TEXT
+                text_block_removed.append(text_block)
+    # 去掉table上的文字block
+    for table_box in tables:
+        for text_block in text_raw_blocks:
+            text_bbox = text_block["bbox"]
+            if _is_in(text_bbox, table_box):
+                text_block['tag'] = ON_TABLE_TEXT
+                text_block_removed.append(text_block)
+
+    for text_block in text_block_removed:
+        if text_block in text_raw_blocks:
+            text_raw_blocks.remove(text_block)
+
+    # 第一步去掉在图片上出现的公式box
+    temp = []
+    for image_box in images:
+        for eq1 in interline_equations:
+            if _is_in_or_part_overlap(image_box, eq1[:4]):
+                temp.append(eq1)
+        for eq2 in inline_equations:
+            if _is_in_or_part_overlap(image_box, eq2[:4]):
+                temp.append(eq2)
+
+    for eq in temp:
+        if eq in interline_equations:
+            interline_equations.remove(eq)
+        if eq in inline_equations:
+            inline_equations.remove(eq)
+
+    # 第二步去掉在表格上出现的公式box
+    temp = []
+    for table_box in tables:
+        for eq1 in interline_equations:
+            if _is_in_or_part_overlap(table_box, eq1[:4]):
+                temp.append(eq1)
+        for eq2 in inline_equations:
+            if _is_in_or_part_overlap(table_box, eq2[:4]):
+                temp.append(eq2)
+
+    for eq in temp:
+        if eq in interline_equations:
+            interline_equations.remove(eq)
+        if eq in inline_equations:
+            inline_equations.remove(eq)
+
+    # 图片和文字重叠，丢掉图片
+    for image_box in images:
+        for text_block in text_raw_blocks:
+            text_bbox = text_block["bbox"]
+            if _is_in_or_part_overlap(image_box, text_bbox):
+                images_backup.append(image_box)
+                break
+    for image_box in images_backup:
+        images.remove(image_box)
+
+    # 图片和图片重叠，两张都暂时不参与版面计算
+    images_dup_index = []
+    for i in range(len(images)):
+        for j in range(i + 1, len(images)):
+            if _is_in_or_part_overlap(images[i], images[j]):
+                images_dup_index.append(i)
+                images_dup_index.append(j)
+
+    dup_idx = set(images_dup_index)
+    for img_id in dup_idx:
+        images_backup.append(images[img_id])
+        images[img_id] = None
+
+    images = [img for img in images if img is not None]
+
+    # 如果行间公式和文字block重叠，放到临时的数据里，防止这些文字box影响到layout计算。通过计算IOU合并行间公式和文字block
+    # 对于这样的文本块删除，然后保留行间公式的大小不变。
+    # 当计算完毕layout，这部分再合并回来
+    text_block_removed_2 = []
+    # for text_block in text_raw_blocks:
+    #     text_bbox = text_block["bbox"]
+    #     for eq in interline_equations:
+    #         ratio = calculate_overlap_area_2_minbox_area_ratio(text_bbox, eq[:4])
+    #         if ratio>0.05:
+    #             text_block['tag'] = "belong-to-interline-equation"
+    #             text_block_removed_2.append(text_block)
+    #             break
+
+    # for tb in text_block_removed_2:
+    #     if tb in text_raw_blocks:
+    #         text_raw_blocks.remove(tb)
+
+    # text_block_removed = text_block_removed + text_block_removed_2
+
+    return images, tables, interline_equations, inline_equations, text_raw_blocks, text_block_removed, images_backup, text_block_removed_2
+
+
+def check_text_block_horizontal_overlap(text_blocks: list, header, footer) -> bool:
+    """
+    检查文本block之间的水平重叠情况，这种情况如果发生，那么这个pdf就不再继续处理了。
+    因为这种情况大概率发生了公式没有被检测出来。
+    
+    """
+    if len(text_blocks) == 0:
+        return False
+
+    page_min_y = 0
+    page_max_y = max(yy['bbox'][3] for yy in text_blocks)
+
+    def __max_y(lst: list):
+        if len(lst) > 0:
+            return max([item[1] for item in lst])
+        return page_min_y
+
+    def __min_y(lst: list):
+        if len(lst) > 0:
+            return min([item[3] for item in lst])
+        return page_max_y
+
+    clip_y0 = __max_y(header)
+    clip_y1 = __min_y(footer)
+
+    txt_bboxes = []
+    for text_block in text_blocks:
+        bbox = text_block["bbox"]
+        if bbox[1] >= clip_y0 and bbox[3] <= clip_y1:
+            txt_bboxes.append(bbox)
+
+    for i in range(len(txt_bboxes)):
+        for j in range(i + 1, len(txt_bboxes)):
+            if _is_left_overlap(txt_bboxes[i], txt_bboxes[j]) or _is_left_overlap(txt_bboxes[j], txt_bboxes[i]):
+                return True
+
+    return False
+
+
+def check_useful_block_horizontal_overlap(useful_blocks: list) -> bool:
+    """
+    检查文本block之间的水平重叠情况，这种情况如果发生，那么这个pdf就不再继续处理了。
+    因为这种情况大概率发生了公式没有被检测出来。
+
+    """
+    if len(useful_blocks) == 0:
+        return False
+
+    page_min_y = 0
+    page_max_y = max(yy['bbox'][3] for yy in useful_blocks)
+
+    useful_bboxes = []
+    for text_block in useful_blocks:
+        bbox = text_block["bbox"]
+        if bbox[1] >= page_min_y and bbox[3] <= page_max_y:
+            useful_bboxes.append(bbox)
+
+    for i in range(len(useful_bboxes)):
+        for j in range(i + 1, len(useful_bboxes)):
+            area_i = (useful_bboxes[i][2] - useful_bboxes[i][0]) * (useful_bboxes[i][3] - useful_bboxes[i][1])
+            area_j = (useful_bboxes[j][2] - useful_bboxes[j][0]) * (useful_bboxes[j][3] - useful_bboxes[j][1])
+            if _is_left_overlap(useful_bboxes[i], useful_bboxes[j]) or _is_left_overlap(useful_bboxes[j], useful_bboxes[i]):
+                if area_i > area_j:
+                    return True, useful_bboxes[j], useful_bboxes[i]
+                else:
+                    return True, useful_bboxes[i], useful_bboxes[j]
+
+    return False, None, None

magic_pdf/pre_proc/solve_line_alien.py

@@ -0,0 +1,29 @@
+def solve_inline_too_large_interval(pdf_info_dict: dict) -> dict:  # text_block -> json中的preproc_block
+    """解决行内文本间距过大问题"""
+    for i in range(len(pdf_info_dict)):
+
+        text_blocks = pdf_info_dict[f'page_{i}']['preproc_blocks']
+
+        for block in text_blocks:
+
+            x_pre_1, y_pre_1, x_pre_2, y_pre_2 = 0, 0, 0, 0
+            
+            for line in block['lines']:
+
+                x_cur_1, y_cur_1, x_cur_2, y_cur_2 = line['bbox']
+                # line_box = [x1, y1, x2, y2] 
+                if int(y_cur_1) == int(y_pre_1) and int(y_cur_2) == int(y_pre_2):
+                    # if len(line['spans']) == 1:
+                    line['spans'][0]['text'] = ' ' + line['spans'][0]['text']
+                
+                x_pre_1, y_pre_1, x_pre_2, y_pre_2 = line['bbox'] 
+
+    return pdf_info_dict
+
+
+
+
+
+
+
+

magic_pdf/pre_proc/statistics.py

@@ -0,0 +1,12 @@
+
+"""
+统计处需要跨页、全局性的数据
+- 统计出字号从大到小
+- 正文区域占比最高的前5
+- 正文平均行间距
+- 正文平均字间距
+- 正文平均字符宽度
+- 正文平均字符高度
+
+"""
+

magic_pdf/rw/AbsReaderWriter.py

@@ -0,0 +1,34 @@
+from abc import ABC, abstractmethod
+
+
+class AbsReaderWriter(ABC):
+    """
+    同时支持二进制和文本读写的抽象类
+    """
+    MODE_TXT = "text"
+    MODE_BIN = "binary"
+
+    def __init__(self, parent_path):
+        # 初始化代码可以在这里添加，如果需要的话
+        self.parent_path = parent_path # 对于本地目录是父目录，对于s3是会写到这个path下。
+
+    @abstractmethod
+    def read(self, path: str, mode=MODE_TXT):
+        """
+        无论对于本地还是s3的路径，检查如果path是绝对路径，那么就不再 拼接parent_path, 如果是相对路径就拼接parent_path
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def write(self, content: str, path: str, mode=MODE_TXT):
+        """
+        无论对于本地还是s3的路径，检查如果path是绝对路径，那么就不再 拼接parent_path, 如果是相对路径就拼接parent_path
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def read_jsonl(self, path: str, byte_start=0, byte_end=None, encoding='utf-8'):
+        """
+        无论对于本地还是s3的路径，检查如果path是绝对路径，那么就不再 拼接parent_path, 如果是相对路径就拼接parent_path
+        """
+        raise NotImplementedError

magic_pdf/rw/DiskReaderWriter.py

@@ -0,0 +1,66 @@
+import os
+from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
+from loguru import logger
+
+
+MODE_TXT = "text"
+MODE_BIN = "binary"
+
+
+class DiskReaderWriter(AbsReaderWriter):
+
+    def __init__(self, parent_path, encoding="utf-8"):
+        self.path = parent_path
+        self.encoding = encoding
+
+    def read(self, path, mode=MODE_TXT):
+        if os.path.isabs(path):
+            abspath = path
+        else:
+            abspath = os.path.join(self.path, path)
+        if not os.path.exists(abspath):
+            logger.error(f"文件 {abspath} 不存在")
+            raise Exception(f"文件 {abspath} 不存在")
+        if mode == MODE_TXT:
+            with open(abspath, "r", encoding=self.encoding) as f:
+                return f.read()
+        elif mode == MODE_BIN:
+            with open(abspath, "rb") as f:
+                return f.read()
+        else:
+            raise ValueError("Invalid mode. Use 'text' or 'binary'.")
+
+    def write(self, content, path, mode=MODE_TXT):
+        if os.path.isabs(path):
+            abspath = path
+        else:
+            abspath = os.path.join(self.path, path)
+        directory_path = os.path.dirname(abspath)
+        if not os.path.exists(directory_path):
+            os.makedirs(directory_path)
+        if mode == MODE_TXT:
+            with open(abspath, "w", encoding=self.encoding, errors="replace") as f:
+                f.write(content)
+
+        elif mode == MODE_BIN:
+            with open(abspath, "wb") as f:
+                f.write(content)
+        else:
+            raise ValueError("Invalid mode. Use 'text' or 'binary'.")
+
+    def read_jsonl(self, path: str, byte_start=0, byte_end=None, encoding="utf-8"):
+        return self.read(path)
+
+
+# 使用示例
+if __name__ == "__main__":
+    file_path = "io/test/example.txt"
+    drw = DiskReaderWriter("D:\projects\papayfork\Magic-PDF\magic_pdf")
+
+    # 写入内容到文件
+    drw.write(b"Hello, World!", path="io/test/example.txt", mode="binary")
+
+    # 从文件读取内容
+    content = drw.read(path=file_path)
+    if content:
+        logger.info(f"从 {file_path} 读取的内容: {content}")

magic_pdf/rw/S3ReaderWriter.py

@@ -0,0 +1,107 @@
+from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
+from magic_pdf.libs.commons import parse_aws_param, parse_bucket_key, join_path
+import boto3
+from loguru import logger
+from boto3.s3.transfer import TransferConfig
+from botocore.config import Config
+import os
+
+MODE_TXT = "text"
+MODE_BIN = "binary"
+
+
+class S3ReaderWriter(AbsReaderWriter):
+    def __init__(self, ak: str, sk: str, endpoint_url: str, addressing_style: str = 'auto', parent_path: str = ''):
+        self.client = self._get_client(ak, sk, endpoint_url, addressing_style)
+        self.path = parent_path
+
+    def _get_client(self, ak: str, sk: str, endpoint_url: str, addressing_style: str):
+        s3_client = boto3.client(
+            service_name="s3",
+            aws_access_key_id=ak,
+            aws_secret_access_key=sk,
+            endpoint_url=endpoint_url,
+            config=Config(s3={"addressing_style": addressing_style},
+                          retries={'max_attempts': 5, 'mode': 'standard'}),
+        )
+        return s3_client
+
+    def read(self, s3_relative_path, mode=MODE_TXT, encoding="utf-8"):
+        if s3_relative_path.startswith("s3://"):
+            s3_path = s3_relative_path
+        else:
+            s3_path = join_path(self.path, s3_relative_path)
+        bucket_name, key = parse_bucket_key(s3_path)
+        res = self.client.get_object(Bucket=bucket_name, Key=key)
+        body = res["Body"].read()
+        if mode == MODE_TXT:
+            data = body.decode(encoding)  # Decode bytes to text
+        elif mode == MODE_BIN:
+            data = body
+        else:
+            raise ValueError("Invalid mode. Use 'text' or 'binary'.")
+        return data
+
+    def write(self, content, s3_relative_path, mode=MODE_TXT, encoding="utf-8"):
+        if s3_relative_path.startswith("s3://"):
+            s3_path = s3_relative_path
+        else:
+            s3_path = join_path(self.path, s3_relative_path)
+        if mode == MODE_TXT:
+            body = content.encode(encoding)  # Encode text data as bytes
+        elif mode == MODE_BIN:
+            body = content
+        else:
+            raise ValueError("Invalid mode. Use 'text' or 'binary'.")
+        bucket_name, key = parse_bucket_key(s3_path)
+        self.client.put_object(Body=body, Bucket=bucket_name, Key=key)
+        logger.info(f"内容已写入 {s3_path} ")
+
+    def read_jsonl(self, path: str, byte_start=0, byte_end=None, mode=MODE_TXT, encoding='utf-8'):
+        if path.startswith("s3://"):
+            s3_path = path
+        else:
+            s3_path = join_path(self.path, path)
+        bucket_name, key = parse_bucket_key(s3_path)
+
+        range_header = f'bytes={byte_start}-{byte_end}' if byte_end else f'bytes={byte_start}-'
+        res = self.client.get_object(Bucket=bucket_name, Key=key, Range=range_header)
+        body = res["Body"].read()
+        if mode == MODE_TXT:
+            data = body.decode(encoding)  # Decode bytes to text
+        elif mode == MODE_BIN:
+            data = body
+        else:
+            raise ValueError("Invalid mode. Use 'text' or 'binary'.")
+        return data
+
+
+if __name__ == "__main__":
+    # Config the connection info
+    ak = ""
+    sk = ""
+    endpoint_url = ""
+    addressing_style = "auto"
+    bucket_name = ""
+    # Create an S3ReaderWriter object
+    s3_reader_writer = S3ReaderWriter(ak, sk, endpoint_url, addressing_style, "s3://bucket_name/")
+
+    # Write text data to S3
+    text_data = "This is some text data"
+    s3_reader_writer.write(data=text_data, s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=MODE_TXT)
+
+    # Read text data from S3
+    text_data_read = s3_reader_writer.read(s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=MODE_TXT)
+    logger.info(f"Read text data from S3: {text_data_read}")
+    # Write binary data to S3
+    binary_data = b"This is some binary data"
+    s3_reader_writer.write(data=text_data, s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=MODE_BIN)
+
+    # Read binary data from S3
+    binary_data_read = s3_reader_writer.read(s3_relative_path=f"s3://{bucket_name}/ebook/test/test.json", mode=MODE_BIN)
+    logger.info(f"Read binary data from S3: {binary_data_read}")
+
+    # Range Read text data from S3
+    binary_data_read = s3_reader_writer.read_jsonl(path=f"s3://{bucket_name}/ebook/test/test.json",
+                                                   byte_start=0, byte_end=10, mode=MODE_BIN)
+    logger.info(f"Read binary data from S3: {binary_data_read}")

magic_pdf/spark/spark_api.py

@@ -0,0 +1,51 @@
+from loguru import logger
+
+from magic_pdf.libs.drop_reason import DropReason
+
+
+def get_data_source(jso: dict):
+    data_source = jso.get("data_source")
+    if data_source is None:
+        data_source = jso.get("file_source")
+    return data_source
+
+
+def get_data_type(jso: dict):
+    data_type = jso.get("data_type")
+    if data_type is None:
+        data_type = jso.get("file_type")
+    return data_type
+
+
+def get_bookid(jso: dict):
+    book_id = jso.get("bookid")
+    if book_id is None:
+        book_id = jso.get("original_file_id")
+    return book_id
+
+
+def exception_handler(jso: dict, e):
+    logger.exception(e)
+    jso["_need_drop"] = True
+    jso["_drop_reason"] = DropReason.Exception
+    jso["_exception"] = f"ERROR: {e}"
+    return jso
+
+
+def get_bookname(jso: dict):
+    data_source = get_data_source(jso)
+    file_id = jso.get("file_id")
+    book_name = f"{data_source}/{file_id}"
+    return book_name
+
+
+def spark_json_extractor(jso: dict) -> dict:
+
+    """
+    从json中提取数据，返回一个dict
+    """
+
+    return {
+        "_pdf_type": jso["_pdf_type"],
+        "model_list": jso["doc_layout_result"],
+    }