瀏覽代碼

resume download of models from directory

Bruce MacDonald 1 年之前
父節點
當前提交
a86a4f8c26
共有 2 個文件被更改,包括 28 次插入12 次删除
  1. 12 6
      ollama/cmd/cli.py
  2. 16 6
      ollama/model.py

+ 12 - 6
ollama/cmd/cli.py

@@ -175,12 +175,18 @@ def search(*args, **kwargs):
 
 
 def pull(*args, **kwargs):
-    model.pull(model_name=kwargs.pop('model'), *args, **kwargs)
-    print("Up to date.")
+    try:
+        model.pull(model_name=kwargs.pop('model'), *args, **kwargs)
+        print("Up to date.")
+    except Exception as e:
+        print(f"An error occurred: {e}")
 
 
 def run(*args, **kwargs):
-    name = model.pull(model_name=kwargs.pop('model'), *args, **kwargs)
-    kwargs.update({"model": name})
-    print(f"Running {name}...")
-    generate(*args, **kwargs)
+    try:
+        name = model.pull(model_name=kwargs.pop('model'), *args, **kwargs)
+        kwargs.update({"model": name})
+        print(f"Running {name}...")
+        generate(*args, **kwargs)
+    except Exception as e:
+        print(f"An error occurred: {e}")

+ 16 - 6
ollama/model.py

@@ -123,16 +123,26 @@ def download_file(download_url, file_name, file_size):
 
 
 def pull(model_name, *args, **kwargs):
-    maybe_existing_model_location = MODELS_CACHE_PATH / str(model_name + '.bin')
-    if path.exists(model_name) or path.exists(maybe_existing_model_location):
-        # a file on the filesystem is being specified
-        return model_name
     # check the remote model location and see if it needs to be downloaded
     url = model_name
     file_name = ""
     if not validators.url(url) and not url.startswith('huggingface.co'):
-        url = get_url_from_directory(model_name)
-        file_name = model_name
+        try:
+            url = get_url_from_directory(model_name)
+        except Exception as e:
+            # may not have been able to check remote directory, return now
+            return model_name
+        if url is model_name:
+            # this is not a model from our directory, so can't check remote
+            maybe_existing_model_location = MODELS_CACHE_PATH / str(model_name + '.bin')
+            if path.exists(model_name) or path.exists(maybe_existing_model_location):
+                # a file on the filesystem is being specified
+                return model_name
+            raise Exception("unknown model")
+        else:
+            # this is a model from our directory, check remote
+            file_name = model_name
+
     if not (url.startswith('http://') or url.startswith('https://')):
         url = f'https://{url}'