From 3effcfc96aa8792ad2a6d12fd6e30b95ab1334ea Mon Sep 17 00:00:00 2001 From: miguel Date: Wed, 31 May 2023 11:13:30 +0100 Subject: [PATCH] bug fix from refactoring and pre-commits Signed-off-by: miguel --- deepsearch/model/cli/main.py | 4 +++- .../model/model_download/artifact_management_typer.py | 8 +++++--- deepsearch/model/model_download/artifact_manager.py | 6 +++--- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/deepsearch/model/cli/main.py b/deepsearch/model/cli/main.py index 4ebac174..9aa6363b 100644 --- a/deepsearch/model/cli/main.py +++ b/deepsearch/model/cli/main.py @@ -1,6 +1,8 @@ import typer -from deepsearch.model.model_download.artifact_management_typer import artifact_download_app +from deepsearch.model.model_download.artifact_management_typer import ( + artifact_download_app, +) app = typer.Typer(no_args_is_help=True, add_completion=False) app.add_typer( diff --git a/deepsearch/model/model_download/artifact_management_typer.py b/deepsearch/model/model_download/artifact_management_typer.py index d1b81e48..c52915e6 100644 --- a/deepsearch/model/model_download/artifact_management_typer.py +++ b/deepsearch/model/model_download/artifact_management_typer.py @@ -12,7 +12,7 @@ @artifact_download_app.command() def download( - artifact: str = typer.Option(None, "--artifact", "-a"), + artifact_name: str = typer.Option(None, "--artifact", "-a"), list_artifacts: bool = typer.Option(False, "--list", "-l"), ): @@ -20,8 +20,10 @@ def download( artifacts = artifact_manager.get_index_artifact_list() for artifact in artifacts: typer.echo(artifact["folder_name"]) - elif artifact is not None: - artifact_manager.download_artifact_to_cache(artifact, with_progess_bar=True) + elif artifact_name is not None: + artifact_manager.download_artifact_to_cache( + artifact_name, with_progess_bar=True + ) @artifact_download_app.command() diff --git a/deepsearch/model/model_download/artifact_manager.py b/deepsearch/model/model_download/artifact_manager.py index 5cb4c985..16367d8d 100644 --- a/deepsearch/model/model_download/artifact_manager.py +++ b/deepsearch/model/model_download/artifact_manager.py @@ -34,10 +34,10 @@ def get_artifact_location_in_cache(self, artifact_name: str): if "folder_name" in artifact and artifact["folder_name"] == artifact_name: return artifact - def delete_artifact_from_cache(self, artifact: str): + def delete_artifact_from_cache(self, artifact_name: str): target_artifacts = [] for artifact in self.get_artifact_cache_list(): - if "folder_name" in artifact and artifact["folder_name"] == artifact: + if "folder_name" in artifact and artifact["folder_name"] == artifact_name: target_artifacts.append(artifact) for artifact in target_artifacts: @@ -119,7 +119,7 @@ def _download_file( self, artifact_info: Dict, directory: Any, with_progress_bar: bool = False ) -> str: # Get the filename from the URL - filename = artifact_info["artifact_filename"] + filename = artifact_info["model_filename"] file_path = directory.name + f"/{filename}" # Download the file