Merge pull request #56 from opendatalab/dev

Dev

Author: myhloli

mineru_vl_utils/mineru_client.py

@@ -1,10 +1,12 @@
 import asyncio
 import math
+import os
 import re
 from concurrent.futures import Executor
 from typing import Literal, Sequence
 
 from PIL import Image
+from loguru import logger
 
 from .post_process import post_process
 from .structs import BLOCK_TYPES, ContentBlock
@@ -322,6 +324,15 @@ def __init__(
         max_retries: int = 3,  # for http-client backend only
         retry_backoff_factor: float = 0.5,  # for http-client backend only
     ) -> None:
+        env_debug_value = os.getenv('MINERU_VL_DEBUG_ENABLE', '')
+        if env_debug_value:
+            if env_debug_value.lower() in ['true', '1', 'yes']:
+                debug = True
+            elif env_debug_value.lower() in ['false', '0', 'no']:
+                debug = False
+            else:
+                logger.warning(f'unknown MINERU_VL_DEBUG_ENABLE config: {env_debug_value}, pass')
+
         if backend == "transformers":
             if model is None or processor is None:
                 if not model_path:

mineru_vl_utils/vlm_client/http_client.py

@@ -2,6 +2,7 @@
 import json
 import os
 import re
+from enum import Enum
 from typing import AsyncIterable, Iterable, Sequence
 
 import httpx
@@ -35,6 +36,16 @@ def _get_env(key: str, default: str | None = None) -> str:
     raise ValueError(f"Environment variable {key} is not set.")
 
 
+class HTTPMethod(str, Enum):
+    HEAD = "HEAD"
+    GET = "GET"
+    PUT = "PUT"
+    DELETE = "DELETE"
+    OPTIONS = "OPTIONS"
+    TRACE = "TRACE"
+    POST = "POST"
+
+
 class HttpVlmClient(VlmClient):
     def __init__(
         self,
@@ -120,6 +131,7 @@ def _new_client(self) -> httpx.Client:
                 retry=Retry(
                     total=self.max_retries,
                     backoff_factor=self.retry_backoff_factor,
+                    allowed_methods=list(HTTPMethod),
                 ),
                 transport=httpx.HTTPTransport(
                     limits=httpx.Limits(
@@ -144,6 +156,7 @@ async def _new_aio_client(self) -> httpx.AsyncClient:
                 retry=Retry(
                     total=self.max_retries,
                     backoff_factor=self.retry_backoff_factor,
+                    allowed_methods=list(HTTPMethod),
                 ),
                 transport=httpx.AsyncHTTPTransport(
                     limits=httpx.Limits(