Upload download_models.py

download_models.py

@@ -6,8 +6,9 @@ from modelscope import snapshot_download
 
 
 def download_json(url):
+    # 下载JSON文件
     response = requests.get(url)
-    response.raise_for_status()
+    response.raise_for_status()  # 检查请求是否成功
     return response.json()
 
 
@@ -30,11 +31,6 @@ def download_and_modify_json(url, local_filename, modifications):
 
 
 if __name__ == '__main__':
-    # 在 Docker 环境中使用 /app 目录
-    base_dir = "/app"
-    models_dir = os.path.join(base_dir, "models")
-    os.makedirs(models_dir, exist_ok=True)
-    
     mineru_patterns = [
         "models/Layout/LayoutLMv3/*",
         "models/Layout/YOLO/*",
@@ -44,18 +40,30 @@ if __name__ == '__main__':
         "models/TabRec/StructEqTable/*",
     ]
     
-    # 下载模型到 Docker 容器中的固定路径
-    model_dir = snapshot_download('opendatalab/PDF-Extract-Kit-1.0', 
-                                cache_dir=models_dir,
-                                allow_patterns=mineru_patterns)
+    # 在 Docker 环境中使用 /app 目录
+    base_dir = "/app"
+    cache_dir = os.path.join(base_dir, "model_cache")
+    os.makedirs(cache_dir, exist_ok=True)
+    
+    # 下载模型，使用与原代码相同的参数
+    try:
+        model_dir = snapshot_download('opendatalab/PDF-Extract-Kit-1.0', 
+                                    cache_dir=cache_dir,
+                                    revision='master')
+        model_dir = os.path.join(model_dir, 'models')
+    except TypeError:
+        # 如果 allow_patterns 不支持，则下载完整模型
+        model_dir = snapshot_download('opendatalab/PDF-Extract-Kit-1.0',
+                                    cache_dir=cache_dir)
+        model_dir = os.path.join(model_dir, 'models')
+    
     layoutreader_model_dir = snapshot_download('ppaanngggg/layoutreader',
-                                             cache_dir=models_dir)
+                                             cache_dir=cache_dir)
     
-    model_dir = os.path.join(model_dir, 'models')
     print(f'model_dir is: {model_dir}')
     print(f'layoutreader_model_dir is: {layoutreader_model_dir}')
 
-    # 配置文件保存到 Docker 容器的固定路径
+    # 配置文件保存到 Docker 容器的应用目录
     json_url = 'https://gitee.com/myhloli/MinerU/raw/master/magic-pdf.template.json'
     config_file_name = 'magic-pdf.json'
     config_file = os.path.join(base_dir, config_file_name)