Aistudio 使用 SDK 进行下载 (#10678)

* aistudio SDK

* test: transformer download util

* style: pre-commit check

* test

* fix: remove redundant parameter

* Update aistudio-sdk requirement

Author: Yzc216

paddlenlp/transformers/aistudio_utils.py

@@ -14,7 +14,8 @@
 
 from typing import Optional
 
-from aistudio_sdk.hub import download
+from aistudio_sdk.file_download import model_file_download as download
+from requests import HTTPError
 
 
 class UnauthorizedError(Exception):
@@ -46,22 +47,24 @@ def aistudio_download(
     if revision is not None:
         download_kwargs["revision"] = revision
     if cache_dir is not None:
-        download_kwargs["cache_dir"] = cache_dir
-    res = download(
-        repo_id=repo_id,
-        filename=filename,
-        **download_kwargs,
-    )
-    if "path" in res:
-        return res["path"]
-    else:
-        if res["error_code"] == 10001:
-            raise ValueError("Illegal argument error")
-        elif res["error_code"] == 10002:
-            raise UnauthorizedError(
-                "Unauthorized Access. Please ensure that you have provided the AIStudio Access Token and you have access to the requested asset"
-            )
-        elif res["error_code"] == 12001:
-            raise EntryNotFoundError(f"Cannot find the requested file '{filename}' in repo '{repo_id}'")
-        else:
-            raise Exception(f"Unknown error: {res}")
+        download_kwargs["local_dir"] = cache_dir
+
+    try:
+        return download(
+            repo_id=repo_id,
+            file_path=filename,
+            **download_kwargs,
+        )
+    except ValueError:
+        raise EnvironmentError(
+            f"Cannot find {filename} in the cached files and it looks like {repo_id} is not the path to a directory containing the {filename} or"
+            " \nCheckout your internet connection or see how to run the library in offline mode."
+        )
+    except EntryNotFoundError:
+        raise EnvironmentError(
+            f"Cannot find the requested file {filename} in {repo_id}, please make sure the {filename} under the repo {repo_id}"
+        )
+    except HTTPError as err:
+        raise EnvironmentError(f"There was a specific connection error when trying to load {repo_id}:\n{err}")
+    except Exception:
+        raise EnvironmentError(f"Please make sure the {filename} under the repo {repo_id}")

paddlenlp/utils/download/__init__.py

@@ -30,11 +30,6 @@
 from paddle import __version__
 from requests import HTTPError
 
-from .aistudio_hub_download import (
-    aistudio_hub_download,
-    aistudio_hub_file_exists,
-    aistudio_hub_try_to_load_from_cache,
-)
 from .bos_download import bos_download, bos_file_exists, bos_try_to_load_from_cache
 
 
@@ -146,8 +141,8 @@ def resolve_file_path(
 
     # check cache
     for filename in filenames:
-        cache_file_name = bos_aistudio_hf_try_to_load_from_cache(
-            repo_id, filename, cache_dir, subfolder, revision, repo_type, from_bos, from_aistudio, from_hf_hub
+        cache_file_name = bos_hf_try_to_load_from_cache(
+            repo_id, filename, cache_dir, subfolder, revision, repo_type, from_bos, from_hf_hub
         )
         if from_hf_hub and cache_file_name is _CACHED_NO_EXIST:
             cache_file_name = None
@@ -187,32 +182,25 @@ def resolve_file_path(
                         return None
 
         elif from_aistudio:
-            log_endpoint = "Aistudio Hub"
-            for filename in filenames:
-                download_kwargs["filename"] = filename
-                is_available = bos_aistudio_hf_file_exist(
-                    repo_id,
-                    filename,
-                    subfolder=subfolder,
-                    repo_type=repo_type,
-                    revision=revision,
-                    token=token,
-                    endpoint=endpoint,
-                    from_bos=from_bos,
-                    from_aistudio=from_aistudio,
-                    from_hf_hub=from_hf_hub,
-                )
-                if is_available:
-                    cached_file = aistudio_hub_download(
-                        **download_kwargs,
+            for index, filename in enumerate(filenames):
+                try:
+                    from aistudio_sdk.file_download import (
+                        model_file_download as aistudio_download,
                     )
-                    if cached_file is not None:
-                        return cached_file
+
+                    return aistudio_download(repo_id, filename, revision, local_files_only, local_dir)
+                except Exception:
+                    if index < len(filenames) - 1:
+                        continue
+                    else:
+                        print(f"please make sure one of the {filenames} under the repo {repo_id}")
+                        return None
+
         elif from_hf_hub:
             log_endpoint = "Huggingface Hub"
             for filename in filenames:
                 download_kwargs["filename"] = filename
-                is_available = bos_aistudio_hf_file_exist(
+                is_available = bos_hf_file_exist(
                     repo_id,
                     filename,
                     subfolder=subfolder,
@@ -221,7 +209,6 @@ def resolve_file_path(
                     token=token,
                     endpoint=endpoint,
                     from_bos=from_bos,
-                    from_aistudio=from_aistudio,
                     from_hf_hub=from_hf_hub,
                 )
                 if is_available:
@@ -235,7 +222,7 @@ def resolve_file_path(
             download_kwargs["url"] = url
             for filename in filenames:
                 download_kwargs["filename"] = filename
-                is_available = bos_aistudio_hf_file_exist(
+                is_available = bos_hf_file_exist(
                     repo_id,
                     filename,
                     subfolder=subfolder,
@@ -244,7 +231,6 @@ def resolve_file_path(
                     token=token,
                     endpoint=endpoint,
                     from_bos=from_bos,
-                    from_aistudio=from_aistudio,
                     from_hf_hub=from_hf_hub,
                 )
                 if is_available:
@@ -291,7 +277,7 @@ def resolve_file_path(
         )
 
 
-def bos_aistudio_hf_file_exist(
+def bos_hf_file_exist(
     repo_id: str,
     filename: str,
     *,
@@ -301,7 +287,6 @@ def bos_aistudio_hf_file_exist(
     token: Optional[str] = None,
     endpoint: Optional[str] = None,
     from_bos: bool = True,
-    from_aistudio: bool = False,
     from_hf_hub: bool = False,
 ):
     assert repo_id is not None, "repo_id cannot be None"
@@ -310,16 +295,7 @@ def bos_aistudio_hf_file_exist(
     if subfolder is None:
         subfolder = ""
     filename = os.path.join(subfolder, filename)
-    if from_aistudio:
-        out = aistudio_hub_file_exists(
-            repo_id=repo_id,
-            filename=filename,
-            repo_type=repo_type,
-            revision=revision,
-            token=token,
-            endpoint=endpoint,
-        )
-    elif from_hf_hub:
+    if from_hf_hub:
         out = hf_hub_file_exists(
             repo_id=repo_id,
             filename=filename,
@@ -339,15 +315,14 @@ def bos_aistudio_hf_file_exist(
     return out
 
 
-def bos_aistudio_hf_try_to_load_from_cache(
+def bos_hf_try_to_load_from_cache(
     repo_id: str,
     filename: str,
     cache_dir: Union[str, Path, None] = None,
     subfolder: str = None,
     revision: Optional[str] = None,
     repo_type: Optional[str] = None,
     from_bos: bool = True,
-    from_aistudio: bool = False,
     from_hf_hub: bool = False,
 ):
     if subfolder is None:
@@ -359,9 +334,7 @@ def bos_aistudio_hf_try_to_load_from_cache(
         revision=revision,
         repo_type=repo_type,
     )
-    if from_aistudio:
-        return aistudio_hub_try_to_load_from_cache(**load_kwargs)
-    elif from_hf_hub:
+    if from_hf_hub:
         return hf_hub_try_to_load_from_cache(**load_kwargs)
     else:
         return bos_try_to_load_from_cache(**load_kwargs)

requirements.txt

@@ -21,7 +21,7 @@ typer
 rich
 safetensors
 fast_dataindex>=0.1.1 ; platform_system == "Linux"
-aistudio-sdk==0.2.6
+aistudio-sdk>=0.3.0
 jinja2
 regex
 numpy<=1.26.4

tests/utils/test_aistudio_download.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+from paddlenlp.transformers.aistudio_utils import aistudio_download
+from paddlenlp.utils.download import resolve_file_path
+
+
+class TestAistudioDownload(unittest.TestCase):
+    def test_aistudio_download(self):
+        # 设置测试数据
+        repo_id = "PaddleNLP/DeepSeek-R1-Distill-Qwen-1.5B"
+        filename = "model.safetensors"
+        revision = "master"
+        local_dir = "./local/model"
+
+        # 调用待测试的函数
+        result = resolve_file_path(
+            repo_id=repo_id,
+            filenames=filename,
+            revision=revision,
+            from_aistudio=True,
+            local_dir=local_dir,
+        )
+
+        # 验证结果
+        print(result)
+        self.assertEqual(result, f"{local_dir}/{filename}")
+
+    def test_aistudio_download_transformer(self):
+        # 设置测试数据
+        repo_id = "PaddleNLP/DeepSeek-R1-Distill-Qwen-1.5B"
+        filename = "model.safetensors"
+        revision = "master"
+        cache_dir = "./local/model"
+
+        # 调用待测试的函数
+        result = aistudio_download(
+            repo_id=repo_id,
+            filename=filename,
+            revision=revision,
+            cache_dir=cache_dir,
+        )
+
+        # 验证结果
+        print(result)
+        self.assertEqual(result, f"{cache_dir}/{filename}")
+
+
+if __name__ == "__main__":
+    unittest.main()