Commit bc0f6932 authored by myhloli's avatar myhloli

feat(model): add model mode selection for PDF analysis

Introduce a new feature that allows users to choose between a "lite" and a "full"
model mode for PDF document analysis. The "lite" mode uses a faster, less
accurate model, while the "full" mode employs a higher-precision model at the
cost of speed. This selection can be made through the CLI or API, providing
flexibility for different use cases.
parent f8f6ba6f
...@@ -178,8 +178,10 @@ def cli(): ...@@ -178,8 +178,10 @@ def cli():
default="auto", default="auto",
) )
@click.option("--inside_model", type=click.BOOL, default=False, help="使用内置模型测试") @click.option("--inside_model", type=click.BOOL, default=False, help="使用内置模型测试")
def json_command(json, method, inside_model): @click.option("--model_mode", type=click.STRING, default="lite", help="内置模型选择。lite: 快速解析,精度较低,full: 高精度解析,速度较慢")
def json_command(json, method, inside_model, model_mode):
model_config.__use_inside_model__ = inside_model model_config.__use_inside_model__ = inside_model
model_config.__model_mode__ = model_mode
if not json.startswith("s3://"): if not json.startswith("s3://"):
logger.error("usage: magic-pdf json-command --json s3://some_bucket/some_path") logger.error("usage: magic-pdf json-command --json s3://some_bucket/some_path")
...@@ -229,8 +231,10 @@ def json_command(json, method, inside_model): ...@@ -229,8 +231,10 @@ def json_command(json, method, inside_model):
default="auto", default="auto",
) )
@click.option("--inside_model", type=click.BOOL, default=False, help="使用内置模型测试") @click.option("--inside_model", type=click.BOOL, default=False, help="使用内置模型测试")
def local_json_command(local_json, method, inside_model): @click.option("--model_mode", type=click.STRING, default="lite", help="内置模型选择。lite: 快速解析,精度较低,full: 高精度解析,速度较慢")
def local_json_command(local_json, method, inside_model, model_mode):
model_config.__use_inside_model__ = inside_model model_config.__use_inside_model__ = inside_model
model_config.__model_mode__ = model_mode
def read_s3_path(s3path): def read_s3_path(s3path):
bucket, key = parse_s3path(s3path) bucket, key = parse_s3path(s3path)
...@@ -281,8 +285,10 @@ def local_json_command(local_json, method, inside_model): ...@@ -281,8 +285,10 @@ def local_json_command(local_json, method, inside_model):
default="auto", default="auto",
) )
@click.option("--inside_model", type=click.BOOL, default=False, help="使用内置模型测试") @click.option("--inside_model", type=click.BOOL, default=False, help="使用内置模型测试")
def pdf_command(pdf, model, method, inside_model): @click.option("--model_mode", type=click.STRING, default="lite", help="内置模型选择。lite: 快速解析,精度较低,full: 高精度解析,速度较慢")
def pdf_command(pdf, model, method, inside_model, model_mode):
model_config.__use_inside_model__ = inside_model model_config.__use_inside_model__ = inside_model
model_config.__model_mode__ = model_mode
def read_fn(path): def read_fn(path):
disk_rw = DiskReaderWriter(os.path.dirname(path)) disk_rw = DiskReaderWriter(os.path.dirname(path))
......
__use_inside_model__ = False __use_inside_model__ = False
__model_mode__ = "lite"
...@@ -46,7 +46,14 @@ def load_images_from_pdf(pdf_bytes: bytes, dpi=200) -> list: ...@@ -46,7 +46,14 @@ def load_images_from_pdf(pdf_bytes: bytes, dpi=200) -> list:
return images return images
def doc_analyze(pdf_bytes: bytes, ocr: bool = False, show_log: bool = False, model=MODEL.PEK): def doc_analyze(pdf_bytes: bytes, ocr: bool = False, show_log: bool = False):
model = None
if model_config.__model_mode__ == "lite":
model = MODEL.Paddle
elif model_config.__model_mode__ == "full":
model = MODEL.PEK
if model_config.__use_inside_model__: if model_config.__use_inside_model__:
model_init_start = time.time() model_init_start = time.time()
if model == MODEL.Paddle: if model == MODEL.Paddle:
......
...@@ -31,6 +31,7 @@ if __name__ == '__main__': ...@@ -31,6 +31,7 @@ if __name__ == '__main__':
extras_require={ extras_require={
"gpu": ["paddleocr", "paddlepaddle-gpu"], "gpu": ["paddleocr", "paddlepaddle-gpu"],
"cpu": ["paddleocr", "paddlepaddle"], "cpu": ["paddleocr", "paddlepaddle"],
"full-cpu": ["unimernet", "matplotlib", "ultralytics", "paddleocr", "paddlepaddle"],
}, },
description="A practical tool for converting PDF to Markdown", # 简短描述 description="A practical tool for converting PDF to Markdown", # 简短描述
long_description=long_description, # 详细描述 long_description=long_description, # 详细描述
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment