mineru 2.5.4__py3-none-any.whl → 2.6.1__py3-none-any.whl

mineru/backend/pipeline/model_init.py

@@ -7,6 +7,7 @@ from .model_list import AtomicModel
 from ...model.layout.doclayoutyolo import DocLayoutYOLOModel
 from ...model.mfd.yolo_v8 import YOLOv8MFDModel
 from ...model.mfr.unimernet.Unimernet import UnimernetModel
+from ...model.mfr.pp_formulanet_plus_m.predict_formula import FormulaRecognizer
 from ...model.ocr.paddleocr2pytorch.pytorch_paddle import PytorchPaddleOCR
 from ...model.ori_cls.paddle_ori_cls import PaddleOrientationClsModel
 from ...model.table.cls.paddle_table_cls import PaddleTableClsModel
@@ -16,6 +17,15 @@ from ...model.table.rec.unet_table.main import UnetTableModel
 from ...utils.enum_class import ModelPath
 from ...utils.models_download_utils import auto_download_and_get_model_root_path
 
+MFR_MODEL = os.getenv('MINERU_FORMULA_CH_SUPPORT', 'False')
+if MFR_MODEL.lower() in ['true', '1', 'yes']:
+    MFR_MODEL = "pp_formulanet_plus_m"
+elif MFR_MODEL.lower() in ['false', '0', 'no']:
+    MFR_MODEL = "unimernet_small"
+else:
+    logger.warning(f"Invalid MINERU_FORMULA_CH_SUPPORT value: {MFR_MODEL}, set to default 'False'")
+    MFR_MODEL = "unimernet_small"
+
 
 def img_orientation_cls_model_init():
     atom_model_manager = AtomModelSingleton()
@@ -68,7 +78,13 @@ def mfd_model_init(weight, device='cpu'):
 
 
 def mfr_model_init(weight_dir, device='cpu'):
-    mfr_model = UnimernetModel(weight_dir, device)
+    if MFR_MODEL == "unimernet_small":
+        mfr_model = UnimernetModel(weight_dir, device)
+    elif MFR_MODEL == "pp_formulanet_plus_m":
+        mfr_model = FormulaRecognizer(weight_dir, device)
+    else:
+        logger.error('MFR model name not allow')
+        exit(1)
     return mfr_model
 
 
@@ -205,11 +221,17 @@ class MineruPipelineModel:
             )
 
             # 初始化公式解析模型
-            mfr_weight_dir = os.path.join(auto_download_and_get_model_root_path(ModelPath.unimernet_small), ModelPath.unimernet_small)
+            if MFR_MODEL == "unimernet_small":
+                mfr_model_path = ModelPath.unimernet_small
+            elif MFR_MODEL == "pp_formulanet_plus_m":
+                mfr_model_path = ModelPath.pp_formulanet_plus_m
+            else:
+                logger.error('MFR model name not allow')
+                exit(1)
 
             self.mfr_model = atom_model_manager.get_atom_model(
                 atom_model_name=AtomicModel.MFR,
-                mfr_weight_dir=mfr_weight_dir,
+                mfr_weight_dir=str(os.path.join(auto_download_and_get_model_root_path(mfr_model_path), mfr_model_path)),
                 device=self.device,
             )
 

mineru/backend/pipeline/model_json_to_middle_json.py

@@ -5,6 +5,7 @@ import time
 from loguru import logger
 from tqdm import tqdm
 
+from mineru.backend.utils import cross_page_table_merge
 from mineru.utils.config_reader import get_device, get_llm_aided_config, get_formula_enable
 from mineru.backend.pipeline.model_init import AtomModelSingleton
 from mineru.backend.pipeline.para_split import para_split
@@ -20,7 +21,6 @@ from mineru.utils.ocr_utils import OcrConfidence
 from mineru.utils.span_block_fix import fill_spans_in_blocks, fix_discarded_block, fix_block_spans
 from mineru.utils.span_pre_proc import remove_outside_spans, remove_overlaps_low_confidence_spans, \
     remove_overlaps_min_spans, txt_spans_extract
-from mineru.utils.table_merge import merge_table
 from mineru.version import __version__
 from mineru.utils.hash_utils import bytes_md5
 
@@ -231,7 +231,7 @@ def result_to_middle_json(model_list, images_list, pdf_doc, image_writer, lang=N
     para_split(middle_json["pdf_info"])
 
     """表格跨页合并"""
-    merge_table(middle_json["pdf_info"])
+    cross_page_table_merge(middle_json["pdf_info"])
 
     """llm优化"""
     llm_aided_config = get_llm_aided_config()

mineru/backend/pipeline/model_list.py

@@ -7,4 +7,3 @@ class AtomicModel:
     WiredTable = "wired_table"
     TableCls = "table_cls"
     ImgOrientationCls = "img_ori_cls"
-

mineru/backend/utils.py

@@ -0,0 +1,24 @@
+import os
+
+from loguru import logger
+
+from mineru.utils.table_merge import merge_table
+
+
+def cross_page_table_merge(pdf_info: list[dict]):
+    """Merge tables that span across multiple pages in a PDF document.
+
+    Args:
+        pdf_info (list[dict]): A list of dictionaries containing information about each page in the PDF.
+
+    Returns:
+        None
+    """
+    is_merge_table = os.getenv('MINERU_TABLE_MERGE_ENABLE', 'true')
+    if is_merge_table.lower() in ['true', '1', 'yes']:
+        merge_table(pdf_info)
+    elif is_merge_table.lower() in ['false', '0', 'no']:
+        pass
+    else:
+        logger.warning(f'unknown MINERU_TABLE_MERGE_ENABLE config: {is_merge_table}, pass')
+        pass

mineru/backend/vlm/model_output_to_middle_json.py

@@ -5,13 +5,13 @@ import cv2
 import numpy as np
 from loguru import logger
 
+from mineru.backend.utils import cross_page_table_merge
 from mineru.backend.vlm.vlm_magic_model import MagicModel
 from mineru.utils.config_reader import get_table_enable, get_llm_aided_config
 from mineru.utils.cut_image import cut_image_and_table
 from mineru.utils.enum_class import ContentType
 from mineru.utils.hash_utils import bytes_md5
 from mineru.utils.pdf_image_tools import get_crop_img
-from mineru.utils.table_merge import merge_table
 from mineru.version import __version__
 
 
@@ -110,7 +110,7 @@ def result_to_middle_json(model_output_blocks_list, images_list, pdf_doc, image_
     """表格跨页合并"""
     table_enable = get_table_enable(os.getenv('MINERU_VLM_TABLE_ENABLE', 'True').lower() == 'true')
     if table_enable:
-        merge_table(middle_json["pdf_info"])
+        cross_page_table_merge(middle_json["pdf_info"])
 
     """llm优化标题分级"""
     if heading_level_import_success:

mineru/backend/vlm/{custom_logits_processors.py → utils.py}

@@ -3,8 +3,11 @@ import os
 from loguru import logger
 from packaging import version
 
+from mineru.utils.config_reader import get_device
+from mineru.utils.model_utils import get_vram
 
-def enable_custom_logits_processors():
+
+def enable_custom_logits_processors() -> bool:
     import torch
     from vllm import __version__ as vllm_version
 
@@ -38,4 +41,35 @@ def enable_custom_logits_processors():
             return False
     else:
         logger.info(f"compute_capability: {compute_capability} >= 8.0 and vllm version: {vllm_version} >= 0.10.1, enable custom_logits_processors")
-        return True
+        return True
+
+
+def set_default_gpu_memory_utilization() -> float:
+    from vllm import __version__ as vllm_version
+    if version.parse(vllm_version) >= version.parse("0.11.0"):
+        return 0.7
+    else:
+        return 0.5
+
+
+def set_default_batch_size() -> int:
+    try:
+        device = get_device()
+        vram = get_vram(device)
+        if vram is not None:
+            gpu_memory = int(os.getenv('MINERU_VIRTUAL_VRAM_SIZE', round(vram)))
+            if gpu_memory >= 16:
+                batch_size = 8
+            elif gpu_memory >= 8:
+                batch_size = 4
+            else:
+                batch_size = 1
+            logger.info(f'gpu_memory: {gpu_memory} GB, batch_size: {batch_size}')
+        else:
+            # Default batch_ratio when VRAM can't be determined
+            batch_size = 1
+            logger.info(f'Could not determine GPU memory, using default batch_ratio: {batch_size}')
+    except Exception as e:
+        logger.warning(f'Error determining VRAM: {e}, using default batch_ratio: 1')
+        batch_size = 1
+    return batch_size

mineru/backend/vlm/vlm_analyze.py

@@ -4,14 +4,13 @@ import time
 
 from loguru import logger
 
-from .custom_logits_processors import enable_custom_logits_processors
+from .utils import enable_custom_logits_processors, set_default_gpu_memory_utilization, set_default_batch_size
 from .model_output_to_middle_json import result_to_middle_json
 from ...data.data_reader_writer import DataWriter
 from mineru.utils.pdf_image_tools import load_images_from_pdf
 from ...utils.config_reader import get_device
 
 from ...utils.enum_class import ImageType
-from ...utils.model_utils import get_vram
 from ...utils.models_download_utils import auto_download_and_get_model_root_path
 
 from mineru_vl_utils import MinerUClient
@@ -41,7 +40,13 @@ class ModelSingleton:
             processor = None
             vllm_llm = None
             vllm_async_llm = None
-            batch_size = 0
+            batch_size = kwargs.get("batch_size", 0)  # for transformers backend only
+            max_concurrency = kwargs.get("max_concurrency", 100)  # for http-client backend only
+            http_timeout = kwargs.get("http_timeout", 600)  # for http-client backend only
+            # 从kwargs中移除这些参数，避免传递给不相关的初始化函数
+            for param in ["batch_size", "max_concurrency", "http_timeout"]:
+                if param in kwargs:
+                    del kwargs[param]
             if backend in ['transformers', 'vllm-engine', "vllm-async-engine"] and not model_path:
                 model_path = auto_download_and_get_model_root_path("/","vlm")
                 if backend == "transformers":
@@ -68,53 +73,41 @@ class ModelSingleton:
                         model_path,
                         use_fast=True,
                     )
-                    try:
-                        vram = get_vram(device)
-                        if vram is not None:
-                            gpu_memory = int(os.getenv('MINERU_VIRTUAL_VRAM_SIZE', round(vram)))
-                            if gpu_memory >= 16:
-                                batch_size = 8
-                            elif gpu_memory >= 8:
-                                batch_size = 4
-                            else:
-                                batch_size = 1
-                            logger.info(f'gpu_memory: {gpu_memory} GB, batch_size: {batch_size}')
-                        else:
-                            # Default batch_ratio when VRAM can't be determined
-                            batch_size = 1
-                            logger.info(f'Could not determine GPU memory, using default batch_ratio: {batch_size}')
-                    except Exception as e:
-                        logger.warning(f'Error determining VRAM: {e}, using default batch_ratio: 1')
-                        batch_size = 1
-                elif backend == "vllm-engine":
-                    try:
-                        import vllm
-                        from mineru_vl_utils import MinerULogitsProcessor
-                    except ImportError:
-                        raise ImportError("Please install vllm to use the vllm-engine backend.")
-                    if "gpu_memory_utilization" not in kwargs:
-                        kwargs["gpu_memory_utilization"] = 0.5
-                    if "model" not in kwargs:
-                        kwargs["model"] = model_path
-                    if enable_custom_logits_processors() and ("logits_processors" not in kwargs):
-                        kwargs["logits_processors"] = [MinerULogitsProcessor]
-                    # 使用kwargs为 vllm初始化参数
-                    vllm_llm = vllm.LLM(**kwargs)
-                elif backend == "vllm-async-engine":
-                    try:
-                        from vllm.engine.arg_utils import AsyncEngineArgs
-                        from vllm.v1.engine.async_llm import AsyncLLM
-                        from mineru_vl_utils import MinerULogitsProcessor
-                    except ImportError:
-                        raise ImportError("Please install vllm to use the vllm-async-engine backend.")
-                    if "gpu_memory_utilization" not in kwargs:
-                        kwargs["gpu_memory_utilization"] = 0.5
-                    if "model" not in kwargs:
-                        kwargs["model"] = model_path
-                    if enable_custom_logits_processors() and ("logits_processors" not in kwargs):
-                        kwargs["logits_processors"] = [MinerULogitsProcessor]
-                    # 使用kwargs为 vllm初始化参数
-                    vllm_async_llm = AsyncLLM.from_engine_args(AsyncEngineArgs(**kwargs))
+                    if batch_size == 0:
+                        batch_size = set_default_batch_size()
+                else:
+                    if os.getenv('OMP_NUM_THREADS') is None:
+                        os.environ["OMP_NUM_THREADS"] = "1"
+
+                    if backend == "vllm-engine":
+                        try:
+                            import vllm
+                            from mineru_vl_utils import MinerULogitsProcessor
+                        except ImportError:
+                            raise ImportError("Please install vllm to use the vllm-engine backend.")
+                        if "gpu_memory_utilization" not in kwargs:
+                            kwargs["gpu_memory_utilization"] = set_default_gpu_memory_utilization()
+                        if "model" not in kwargs:
+                            kwargs["model"] = model_path
+                        if enable_custom_logits_processors() and ("logits_processors" not in kwargs):
+                            kwargs["logits_processors"] = [MinerULogitsProcessor]
+                        # 使用kwargs为 vllm初始化参数
+                        vllm_llm = vllm.LLM(**kwargs)
+                    elif backend == "vllm-async-engine":
+                        try:
+                            from vllm.engine.arg_utils import AsyncEngineArgs
+                            from vllm.v1.engine.async_llm import AsyncLLM
+                            from mineru_vl_utils import MinerULogitsProcessor
+                        except ImportError:
+                            raise ImportError("Please install vllm to use the vllm-async-engine backend.")
+                        if "gpu_memory_utilization" not in kwargs:
+                            kwargs["gpu_memory_utilization"] = set_default_gpu_memory_utilization()
+                        if "model" not in kwargs:
+                            kwargs["model"] = model_path
+                        if enable_custom_logits_processors() and ("logits_processors" not in kwargs):
+                            kwargs["logits_processors"] = [MinerULogitsProcessor]
+                        # 使用kwargs为 vllm初始化参数
+                        vllm_async_llm = AsyncLLM.from_engine_args(AsyncEngineArgs(**kwargs))
             self._models[key] = MinerUClient(
                 backend=backend,
                 model=model,
@@ -123,6 +116,8 @@ class ModelSingleton:
                 vllm_async_llm=vllm_async_llm,
                 server_url=server_url,
                 batch_size=batch_size,
+                max_concurrency=max_concurrency,
+                http_timeout=http_timeout,
             )
             elapsed = round(time.time() - start_time, 2)
             logger.info(f"get {backend} predictor cost: {elapsed}s")

mineru/backend/vlm/vlm_magic_model.py

@@ -361,7 +361,7 @@ def get_type_blocks(blocks, block_type: Literal["image", "table", "code"]):
     return ret
 
 
-def fix_two_layer_blocks(blocks, fix_type: Literal["image", "table", "code"]):
+def fix_two_layer_blocks_back(blocks, fix_type: Literal["image", "table", "code"]):
     need_fix_blocks = get_type_blocks(blocks, fix_type)
     fixed_blocks = []
     not_include_blocks = []
@@ -404,6 +404,160 @@ def fix_two_layer_blocks(blocks, fix_type: Literal["image", "table", "code"]):
     return fixed_blocks, not_include_blocks
 
 
+def fix_two_layer_blocks(blocks, fix_type: Literal["image", "table", "code"]):
+    need_fix_blocks = get_type_blocks(blocks, fix_type)
+    fixed_blocks = []
+    not_include_blocks = []
+    processed_indices = set()
+
+    # 特殊处理表格类型，确保标题在表格前，注脚在表格后
+    if fix_type == "table":
+        # 收集所有不合适的caption和footnote
+        misplaced_captions = []  # 存储(caption, 原始block索引)
+        misplaced_footnotes = []  # 存储(footnote, 原始block索引)
+
+        # 第一步：移除不符合位置要求的caption和footnote
+        for block_idx, block in enumerate(need_fix_blocks):
+            body = block[f"{fix_type}_body"]
+            body_index = body["index"]
+
+            # 检查caption应在body前或同位置
+            valid_captions = []
+            for caption in block[f"{fix_type}_caption_list"]:
+                if caption["index"] <= body_index:
+                    valid_captions.append(caption)
+                else:
+                    misplaced_captions.append((caption, block_idx))
+            block[f"{fix_type}_caption_list"] = valid_captions
+
+            # 检查footnote应在body后或同位置
+            valid_footnotes = []
+            for footnote in block[f"{fix_type}_footnote_list"]:
+                if footnote["index"] >= body_index:
+                    valid_footnotes.append(footnote)
+                else:
+                    misplaced_footnotes.append((footnote, block_idx))
+            block[f"{fix_type}_footnote_list"] = valid_footnotes
+
+        # 第二步：重新分配不合规的caption到合适的body
+        for caption, original_block_idx in misplaced_captions:
+            caption_index = caption["index"]
+            best_block_idx = None
+            min_distance = float('inf')
+
+            # 寻找索引大于等于caption_index的最近body
+            for idx, block in enumerate(need_fix_blocks):
+                body_index = block[f"{fix_type}_body"]["index"]
+                if body_index >= caption_index and idx != original_block_idx:
+                    distance = body_index - caption_index
+                    if distance < min_distance:
+                        min_distance = distance
+                        best_block_idx = idx
+
+            if best_block_idx is not None:
+                # 找到合适的body，添加到对应block的caption_list
+                need_fix_blocks[best_block_idx][f"{fix_type}_caption_list"].append(caption)
+            else:
+                # 没找到合适的body，作为普通block处理
+                not_include_blocks.append(caption)
+
+        # 第三步：重新分配不合规的footnote到合适的body
+        for footnote, original_block_idx in misplaced_footnotes:
+            footnote_index = footnote["index"]
+            best_block_idx = None
+            min_distance = float('inf')
+
+            # 寻找索引小于等于footnote_index的最近body
+            for idx, block in enumerate(need_fix_blocks):
+                body_index = block[f"{fix_type}_body"]["index"]
+                if body_index <= footnote_index and idx != original_block_idx:
+                    distance = footnote_index - body_index
+                    if distance < min_distance:
+                        min_distance = distance
+                        best_block_idx = idx
+
+            if best_block_idx is not None:
+                # 找到合适的body，添加到对应block的footnote_list
+                need_fix_blocks[best_block_idx][f"{fix_type}_footnote_list"].append(footnote)
+            else:
+                # 没找到合适的body，作为普通block处理
+                not_include_blocks.append(footnote)
+
+        # 第四步:将每个block的caption_list和footnote_list中不连续index的元素提出来作为普通block处理
+        for block in need_fix_blocks:
+            caption_list = block[f"{fix_type}_caption_list"]
+            footnote_list = block[f"{fix_type}_footnote_list"]
+            body_index = block[f"{fix_type}_body"]["index"]
+
+            # 处理caption_list (从body往前看,caption在body之前)
+            if caption_list:
+                # 按index降序排列,从最接近body的开始检查
+                caption_list.sort(key=lambda x: x["index"], reverse=True)
+                filtered_captions = [caption_list[0]]
+                for i in range(1, len(caption_list)):
+                    # 检查是否与前一个caption连续(降序所以是-1)
+                    if caption_list[i]["index"] == caption_list[i - 1]["index"] - 1:
+                        filtered_captions.append(caption_list[i])
+                    else:
+                        # 出现gap,后续所有caption都作为普通block
+                        not_include_blocks.extend(caption_list[i:])
+                        break
+                # 恢复升序
+                filtered_captions.reverse()
+                block[f"{fix_type}_caption_list"] = filtered_captions
+
+            # 处理footnote_list (从body往后看,footnote在body之后)
+            if footnote_list:
+                # 按index升序排列,从最接近body的开始检查
+                footnote_list.sort(key=lambda x: x["index"])
+                filtered_footnotes = [footnote_list[0]]
+                for i in range(1, len(footnote_list)):
+                    # 检查是否与前一个footnote连续
+                    if footnote_list[i]["index"] == footnote_list[i - 1]["index"] + 1:
+                        filtered_footnotes.append(footnote_list[i])
+                    else:
+                        # 出现gap,后续所有footnote都作为普通block
+                        not_include_blocks.extend(footnote_list[i:])
+                        break
+                block[f"{fix_type}_footnote_list"] = filtered_footnotes
+
+    # 构建两层结构blocks
+    for block in need_fix_blocks:
+        body = block[f"{fix_type}_body"]
+        caption_list = block[f"{fix_type}_caption_list"]
+        footnote_list = block[f"{fix_type}_footnote_list"]
+
+        body["type"] = f"{fix_type}_body"
+        for caption in caption_list:
+            caption["type"] = f"{fix_type}_caption"
+            processed_indices.add(caption["index"])
+        for footnote in footnote_list:
+            footnote["type"] = f"{fix_type}_footnote"
+            processed_indices.add(footnote["index"])
+
+        processed_indices.add(body["index"])
+
+        two_layer_block = {
+            "type": fix_type,
+            "bbox": body["bbox"],
+            "blocks": [body],
+            "index": body["index"],
+        }
+        two_layer_block["blocks"].extend([*caption_list, *footnote_list])
+        # 对blocks按index排序
+        two_layer_block["blocks"].sort(key=lambda x: x["index"])
+
+        fixed_blocks.append(two_layer_block)
+
+    # 添加未处理的blocks
+    for block in blocks:
+        block.pop("type", None)
+        if block["index"] not in processed_indices and block not in not_include_blocks:
+            not_include_blocks.append(block)
+
+    return fixed_blocks, not_include_blocks
+
+
 def fix_list_blocks(list_blocks, text_blocks, ref_text_blocks):
     for list_block in list_blocks:
         list_block["blocks"] = []

mineru/cli/common.py

@@ -18,7 +18,7 @@ from mineru.backend.vlm.vlm_analyze import doc_analyze as vlm_doc_analyze
 from mineru.backend.vlm.vlm_analyze import aio_doc_analyze as aio_vlm_doc_analyze
 
 pdf_suffixes = ["pdf"]
-image_suffixes = ["png", "jpeg", "jp2", "webp", "gif", "bmp", "jpg"]
+image_suffixes = ["png", "jpeg", "jp2", "webp", "gif", "bmp", "jpg", "tiff"]
 
 
 def read_fn(path):
@@ -44,34 +44,37 @@ def prepare_env(output_dir, pdf_file_name, parse_method):
 
 
 def convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes, start_page_id=0, end_page_id=None):
+    try:
+        # 从字节数据加载PDF
+        pdf = pdfium.PdfDocument(pdf_bytes)
 
-    # 从字节数据加载PDF
-    pdf = pdfium.PdfDocument(pdf_bytes)
-
-    # 确定结束页
-    end_page_id = end_page_id if end_page_id is not None and end_page_id >= 0 else len(pdf) - 1
-    if end_page_id > len(pdf) - 1:
-        logger.warning("end_page_id is out of range, use pdf_docs length")
-        end_page_id = len(pdf) - 1
+        # 确定结束页
+        end_page_id = end_page_id if end_page_id is not None and end_page_id >= 0 else len(pdf) - 1
+        if end_page_id > len(pdf) - 1:
+            logger.warning("end_page_id is out of range, use pdf_docs length")
+            end_page_id = len(pdf) - 1
 
-    # 创建一个新的PDF文档
-    output_pdf = pdfium.PdfDocument.new()
+        # 创建一个新的PDF文档
+        output_pdf = pdfium.PdfDocument.new()
 
-    # 选择要导入的页面索引
-    page_indices = list(range(start_page_id, end_page_id + 1))
+        # 选择要导入的页面索引
+        page_indices = list(range(start_page_id, end_page_id + 1))
 
-    # 从原PDF导入页面到新PDF
-    output_pdf.import_pages(pdf, page_indices)
+        # 从原PDF导入页面到新PDF
+        output_pdf.import_pages(pdf, page_indices)
 
-    # 将新PDF保存到内存缓冲区
-    output_buffer = io.BytesIO()
-    output_pdf.save(output_buffer)
+        # 将新PDF保存到内存缓冲区
+        output_buffer = io.BytesIO()
+        output_pdf.save(output_buffer)
 
-    # 获取字节数据
-    output_bytes = output_buffer.getvalue()
+        # 获取字节数据
+        output_bytes = output_buffer.getvalue()
 
-    pdf.close()  # 关闭原PDF文档以释放资源
-    output_pdf.close()  # 关闭新PDF文档以释放资源
+        pdf.close()  # 关闭原PDF文档以释放资源
+        output_pdf.close()  # 关闭新PDF文档以释放资源
+    except Exception as e:
+        logger.warning(f"Error in converting PDF bytes: {e}, Using original PDF bytes.")
+        output_bytes = pdf_bytes
 
     return output_bytes
 

mineru/cli/fast_api.py

@@ -177,10 +177,7 @@ async def parse_pdf(
                             zf.write(path, arcname=os.path.join(safe_pdf_name, f"{safe_pdf_name}_middle.json"))
 
                     if return_model_output:
-                        if backend.startswith("pipeline"):
-                            path = os.path.join(parse_dir, f"{pdf_name}_model.json")
-                        else:
-                            path = os.path.join(parse_dir, f"{pdf_name}_model_output.txt")
+                        path = os.path.join(parse_dir, f"{pdf_name}_model.json")
                         if os.path.exists(path): 
                             zf.write(path, arcname=os.path.join(safe_pdf_name, os.path.basename(path)))
 
@@ -220,10 +217,7 @@ async def parse_pdf(
                     if return_middle_json:
                         data["middle_json"] = get_infer_result("_middle.json", pdf_name, parse_dir)
                     if return_model_output:
-                        if backend.startswith("pipeline"):
-                            data["model_output"] = get_infer_result("_model.json", pdf_name, parse_dir)
-                        else:
-                            data["model_output"] = get_infer_result("_model_output.txt", pdf_name, parse_dir)
+                        data["model_output"] = get_infer_result("_model.json", pdf_name, parse_dir)
                     if return_content_list:
                         data["content_list"] = get_infer_result("_content_list.json", pdf_name, parse_dir)
                     if return_images: