From 3ce3caaf65bff7014bbd13a67e51ea0ecd0bbc01 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Fri, 30 Jun 2023 11:54:46 -0700 Subject: [PATCH] fix pull model name --- ollama/cmd/cli.py | 2 +- ollama/model.py | 18 +++++++++--------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/ollama/cmd/cli.py b/ollama/cmd/cli.py index 61af0131..00f696de 100644 --- a/ollama/cmd/cli.py +++ b/ollama/cmd/cli.py @@ -151,7 +151,7 @@ def pull(*args, **kwargs): def run(*args, **kwargs): - name = model.pull(*args, **kwargs) + name = model.pull(model_name=kwargs.pop('model'), *args, **kwargs) kwargs.update({"model": name}) print(f"Running {name}...") generate(*args, **kwargs) diff --git a/ollama/model.py b/ollama/model.py index ec59686c..0c02a25d 100644 --- a/ollama/model.py +++ b/ollama/model.py @@ -110,25 +110,25 @@ def download_file(download_url, file_name, file_size): return local_filename -def pull(model, *args, **kwargs): - if path.exists(model): +def pull(model_name, *args, **kwargs): + if path.exists(model_name): # a file on the filesystem is being specified - return model + return model_name # check the remote model location and see if it needs to be downloaded - url = model + url = model_name file_name = "" if not validators.url(url) and not url.startswith('huggingface.co'): - url = get_url_from_directory(model) - file_name = model + url = get_url_from_directory(model_name) + file_name = model_name if not (url.startswith('http://') or url.startswith('https://')): url = f'https://{url}' if not validators.url(url): - if model in models(MODELS_CACHE_PATH): + if model_name in models(MODELS_CACHE_PATH): # the model is already downloaded, and specified by name - return model - raise Exception(f'Unknown model {model}') + return model_name + raise Exception(f'Unknown model {model_name}') local_filename = download_from_repo(url, file_name)