Browse Source

Refactor async function and improve output directory handling in prediction

myhloli 1 tháng trước cách đây
mục cha
commit
44fdeb663f
2 tập tin đã thay đổi với 10 bổ sung10 xóa
  1. 5 5
      projects/multi_gpu_v2/client.py
  2. 5 5
      projects/multi_gpu_v2/server.py

+ 5 - 5
projects/multi_gpu_v2/client.py

@@ -1,11 +1,10 @@
 import base64
-import requests
 import os
 from loguru import logger
 import asyncio
 import aiohttp
 
-async def mineru_parse_async(session, file_path, server_url='http://127.0.0.1:8000/predict', **options):
+async def mineru_parse_async(session, file_path, url='http://127.0.0.1:8000/predict', **options):
     """
     Asynchronous version of the parse function.
     """
@@ -20,7 +19,7 @@ async def mineru_parse_async(session, file_path, server_url='http://127.0.0.1:80
         }
 
         # Use the aiohttp session to send the request
-        async with session.post(server_url, json=payload) as response:
+        async with session.post(url, json=payload) as response:
             if response.status == 200:
                 result = await response.json()
                 logger.info(f"✅ Processed: {file_path} -> {result.get('output_dir', 'N/A')}")
@@ -61,9 +60,10 @@ async def main():
         # === Custom Options ===
         custom_options = {
             'backend': 'pipeline', 'lang': 'ch', 'method': 'auto',
-            'formula_enable': True, 'table_enable': True
+            'formula_enable': True, 'table_enable': True,
+            # Example for remote vlm server (vllm/sglang/lmdeploy...)
+            # 'backend': 'vlm-http-client', 'server_url': 'http://127.0.0.1:30000',
         }
-        # 'backend': 'vlm-vllm-engine' requires 8+ GB VRAM per worker
 
         custom_tasks = [mineru_parse_async(session, file_path, **custom_options) for file_path in existing_files[2:]]
 

+ 5 - 5
projects/multi_gpu_v2/server.py

@@ -61,8 +61,8 @@ class MinerUAPI(ls.LitAPI):
     def predict(self, inputs):
         """Call MinerU's do_parse - same as CLI"""
         input_path = inputs['input_path']
-        output_dir = Path(self.output_dir) / Path(input_path).stem
-        
+        output_dir = Path(self.output_dir)
+
         try:
             os.makedirs(output_dir, exist_ok=True)
             
@@ -82,9 +82,9 @@ class MinerUAPI(ls.LitAPI):
                 start_page_id=inputs['start_page_id'],
                 end_page_id=inputs['end_page_id']
             )
-            
-            return str(output_dir)
-            
+
+            return str(output_dir/Path(input_path).stem)
+
         except Exception as e:
             logger.error(f"Processing failed: {e}")
             raise HTTPException(status_code=500, detail=str(e))