Merge pull request #29 from jmorganca/pull-model-name

Pull model name
This commit is contained in:
Michael Yang 2023-06-30 11:58:58 -07:00 committed by GitHub
commit 9e5dfc66a3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 13 additions and 13 deletions

View file

@ -151,7 +151,7 @@ def pull(*args, **kwargs):
def run(*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}) kwargs.update({"model": name})
print(f"Running {name}...") print(f"Running {name}...")
generate(*args, **kwargs) generate(*args, **kwargs)

View file

@ -30,7 +30,7 @@ def load(model_name, models={}):
if not models.get(model_name, None): if not models.get(model_name, None):
model_path = path.expanduser(model_name) model_path = path.expanduser(model_name)
if not path.exists(model_path): if not path.exists(model_path):
model_path = path.join(MODELS_CACHE_PATH, model_name + ".bin") model_path = MODELS_CACHE_PATH / model_name + ".bin"
runners = { runners = {
model_type: cls model_type: cls

View file

@ -7,7 +7,7 @@ from tqdm import tqdm
MODELS_MANIFEST = 'https://ollama.ai/api/models' MODELS_MANIFEST = 'https://ollama.ai/api/models'
MODELS_CACHE_PATH = path.join(Path.home(), '.ollama', 'models') MODELS_CACHE_PATH = Path.home() / '.ollama' / 'models'
def models(*args, **kwargs): def models(*args, **kwargs):
@ -78,7 +78,7 @@ def find_bin_file(json_response, location, branch):
def download_file(download_url, file_name, file_size): def download_file(download_url, file_name, file_size):
local_filename = path.join(MODELS_CACHE_PATH, file_name) + '.bin' local_filename = MODELS_CACHE_PATH / file_name + '.bin'
first_byte = path.getsize(local_filename) if path.exists(local_filename) else 0 first_byte = path.getsize(local_filename) if path.exists(local_filename) else 0
@ -110,25 +110,25 @@ def download_file(download_url, file_name, file_size):
return local_filename return local_filename
def pull(model, *args, **kwargs): def pull(model_name, *args, **kwargs):
if path.exists(model): if path.exists(model_name):
# a file on the filesystem is being specified # 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 # check the remote model location and see if it needs to be downloaded
url = model url = model_name
file_name = "" file_name = ""
if not validators.url(url) and not url.startswith('huggingface.co'): if not validators.url(url) and not url.startswith('huggingface.co'):
url = get_url_from_directory(model) url = get_url_from_directory(model_name)
file_name = model file_name = model_name
if not (url.startswith('http://') or url.startswith('https://')): if not (url.startswith('http://') or url.startswith('https://')):
url = f'https://{url}' url = f'https://{url}'
if not validators.url(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 # the model is already downloaded, and specified by name
return model return model_name
raise Exception(f'Unknown model {model}') raise Exception(f'Unknown model {model_name}')
local_filename = download_from_repo(url, file_name) local_filename = download_from_repo(url, file_name)