ollama run command

ollama/cmd/cli.py

@@ -29,9 +29,13 @@ def main():
     add_parser.set_defaults(fn=add)
 
     pull_parser = subparsers.add_parser("pull")
-    pull_parser.add_argument("remote")
+    pull_parser.add_argument("model")
     pull_parser.set_defaults(fn=pull)
 
+    pull_parser = subparsers.add_parser("run")
+    pull_parser.add_argument("model")
+    pull_parser.set_defaults(fn=run)
+
     args = parser.parse_args()
     args = vars(args)
 
@@ -52,8 +56,8 @@ def list_models(*args, **kwargs):
 
 
 def generate(*args, **kwargs):
-    if prompt := kwargs.get('prompt'):
-        print('>>>', prompt, flush=True)
+    if prompt := kwargs.get("prompt"):
+        print(">>>", prompt, flush=True)
         generate_oneshot(*args, **kwargs)
         return
 
@@ -79,7 +83,7 @@ def generate_oneshot(*args, **kwargs):
 
 def generate_interactive(*args, **kwargs):
     while True:
-        print('>>> ', end='', flush=True)
+        print(">>> ", end="", flush=True)
         line = next(sys.stdin)
         if not line:
             return
@@ -90,7 +94,7 @@ def generate_interactive(*args, **kwargs):
 
 def generate_batch(*args, **kwargs):
     for line in sys.stdin:
-        print('>>> ', line, end='', flush=True)
+        print(">>> ", line, end="", flush=True)
         kwargs.update({"prompt": line})
         generate_oneshot(*args, **kwargs)
 
@@ -101,3 +105,10 @@ def add(model, models_home):
 
 def pull(*args, **kwargs):
     model.pull(*args, **kwargs)
+
+
+def run(*args, **kwargs):
+    name = model.pull(*args, **kwargs)
+    kwargs.update({"model": name})
+    print(f"Running {name}...")
+    generate(*args, **kwargs)

ollama/model.py

@@ -1,5 +1,6 @@
 import os
 import requests
+import validators
 from urllib.parse import urlsplit, urlunsplit
 from tqdm import tqdm
 
@@ -12,33 +13,36 @@ def models(models_home=".", *args, **kwargs):
                 yield base, os.path.join(root, file)
 
 
-def pull(remote, models_home=".", *args, **kwargs):
-    if not (remote.startswith("http://") or remote.startswith("https://")):
-        remote = f"https://{remote}"
+def pull(model, models_home=".", *args, **kwargs):
+    url = model
+    if not (url.startswith("http://") or url.startswith("https://")):
+        url = f"https://{url}"
 
-    parts = urlsplit(remote)
+    parts = urlsplit(url)
     path_parts = parts.path.split("/tree/")
 
     if len(path_parts) == 1:
-        model = path_parts[0]
+        url = path_parts[0]
         branch = "main"
     else:
-        model, branch = path_parts
+        url, branch = path_parts
 
-    model = model.strip("/")
+    url = url.strip("/")
 
     # Reconstruct the URL
     new_url = urlunsplit(
         (
             "https",
             parts.netloc,
-            f"/api/models/{model}/tree/{branch}",
+            f"/api/models/{url}/tree/{branch}",
             parts.query,
             parts.fragment,
         )
     )
 
-    print(f"Pulling {parts.netloc}/{model}...")
+    if not validators.url(new_url):
+        # this may just be a local model location
+        return model
 
     response = requests.get(new_url)
     response.raise_for_status()  # Raises stored HTTPError, if one occurred
@@ -47,22 +51,28 @@ def pull(remote, models_home=".", *args, **kwargs):
 
     # get the last bin file we find, this is probably the most up to date
     download_url = None
+    file_size = 0
     for file_info in json_response:
         if file_info.get("type") == "file" and file_info.get("path").endswith(".bin"):
             f_path = file_info.get("path")
-            download_url = f"https://huggingface.co/{model}/resolve/{branch}/{f_path}"
+            download_url = f"https://huggingface.co/{url}/resolve/{branch}/{f_path}"
+            file_size = file_info.get("size")
 
     if download_url is None:
         raise Exception("No model found")
 
-    local_filename = os.path.join(models_home, os.path.basename(model)) + ".bin"
+    local_filename = os.path.join(models_home, os.path.basename(url)) + ".bin"
 
     # Check if file already exists
+    first_byte = 0
     if os.path.exists(local_filename):
         # TODO: check if the file is the same SHA
         first_byte = os.path.getsize(local_filename)
-    else:
-        first_byte = 0
+
+    if first_byte >= file_size:
+        return local_filename
+
+    print(f"Pulling {parts.netloc}/{model}...")
 
     # If file size is non-zero, resume download
     if first_byte != 0:
@@ -87,3 +97,5 @@ def pull(remote, models_home=".", *args, **kwargs):
         for data in response.iter_content(chunk_size=1024):
             size = file.write(data)
             bar.update(size)
+
+    return local_filename