From 3f99deb8b9ba6b5b993e19f77e4b95f12a5f55fd Mon Sep 17 00:00:00 2001 From: nghiavt1 Date: Tue, 16 Jul 2024 09:53:43 +0700 Subject: [PATCH] fix(exporters): auth token --- optimum/exporters/onnx/__main__.py | 4 ++-- optimum/exporters/tasks.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/optimum/exporters/onnx/__main__.py b/optimum/exporters/onnx/__main__.py index 1e36af06ad..7f943512de 100644 --- a/optimum/exporters/onnx/__main__.py +++ b/optimum/exporters/onnx/__main__.py @@ -224,9 +224,9 @@ def main_export( original_task = task task = TasksManager.map_from_synonym(task) - framework = TasksManager.determine_framework(model_name_or_path, subfolder=subfolder, framework=framework) + framework = TasksManager.determine_framework(model_name_or_path, subfolder=subfolder, framework=framework, token=token) library_name = TasksManager.infer_library_from_model( - model_name_or_path, subfolder=subfolder, library_name=library_name + model_name_or_path, subfolder=subfolder, library_name=library_name, token=token ) torch_dtype = None diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py index c0221f7bf6..d6c0675197 100644 --- a/optimum/exporters/tasks.py +++ b/optimum/exporters/tasks.py @@ -1457,6 +1457,7 @@ def determine_framework( subfolder: str = "", framework: Optional[str] = None, cache_dir: str = HUGGINGFACE_HUB_CACHE, + token: str = None, ) -> str: """ Determines the framework to use for the export. @@ -1484,7 +1485,7 @@ def determine_framework( if framework is not None: return framework - all_files, request_exception = TasksManager.get_model_files(model_name_or_path, subfolder, cache_dir) + all_files, request_exception = TasksManager.get_model_files(model_name_or_path, subfolder, cache_dir, token=token) pt_weight_name = Path(WEIGHTS_NAME).stem pt_weight_extension = Path(WEIGHTS_NAME).suffix