Add hf:// as an alias to huggingface://

Also check start of string contains :// , by checking if start of the string simply starts with things like "hf", "ollama", etc. we are kinda polluting the namespace and not allowing models to start with these terms. Signed-off-by: Eric Curtin <[email protected]>
containers · Oct 18, 2024 · 25335fc · 25335fc
1 parent 88884d8
commit 25335fc
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 6 deletions.
diff --git a/ramalama/cli.py b/ramalama/cli.py
@@ -193,7 +193,7 @@ def login_parser(subparsers):
 def login_cli(args):
     transport = args.TRANSPORT
     if transport != "":
-        transport = os.getenv("RAMALAMA_TRANSPORT")
+        transport = os.getenv("RAMALAMA_TRANSPORT") + "://"
     model = New(str(transport))
     return model.login(args)
 
@@ -211,7 +211,7 @@ def logout_parser(subparsers):
 def logout_cli(args):
     transport = args.TRANSPORT
     if transport != "":
-        transport = os.getenv("RAMALAMA_TRANSPORT")
+        transport = os.getenv("RAMALAMA_TRANSPORT") + "://"
     model = New(str(transport))
     return model.logout(args)
 
@@ -679,11 +679,11 @@ def dry_run(args):
 
 
 def New(model):
-    if model.startswith("huggingface"):
+    if model.startswith("huggingface://") or model.startswith("hf://"):
         return Huggingface(model)
-    if model.startswith("ollama"):
+    if model.startswith("ollama://"):
         return Ollama(model)
-    if model.startswith("oci") | model.startswith("docker"):
+    if model.startswith("oci://") or model.startswith("docker://"):
         return OCI(model)
 
     transport = os.getenv("RAMALAMA_TRANSPORT")

diff --git a/ramalama/huggingface.py b/ramalama/huggingface.py
@@ -12,7 +12,9 @@
 
 class Huggingface(Model):
     def __init__(self, model):
-        super().__init__(model.removeprefix("huggingface://"))
+        model = model.removeprefix("huggingface://")
+        model = model.removeprefix("hf://")
+        super().__init__(model)
         self.type = "HuggingFace"
         split = self.model.rsplit("/", 1)
         self.directory = ""

diff --git a/test/system/050-pull.bats b/test/system/050-pull.bats
@@ -25,6 +25,11 @@ load helpers
 
 # bats test_tags=distro-integration
 @test "ramalama pull huggingface" {
+    run_ramalama pull hf://afrideva/Tiny-Vicuna-1B-GGUF/tiny-vicuna-1b.q2_k.gguf
+    run_ramalama list
+    is "$output" ".*afrideva/Tiny-Vicuna-1B-GGUF/tiny-vicuna-1b.q2_k" "image was actually pulled locally"
+    run_ramalama rm hf://afrideva/Tiny-Vicuna-1B-GGUF/tiny-vicuna-1b.q2_k.gguf
+
     run_ramalama pull huggingface://afrideva/Tiny-Vicuna-1B-GGUF/tiny-vicuna-1b.q2_k.gguf
     run_ramalama list
     is "$output" ".*afrideva/Tiny-Vicuna-1B-GGUF/tiny-vicuna-1b.q2_k" "image was actually pulled locally"