Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 42 additions & 32 deletions ramalama/url.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,35 @@ def extract_model_identifiers(self):

return model_name, model_tag, model_organization

def _assemble_split_file_list(self, snapshot_hash: str) -> list[SnapshotFile]:
files: list[SnapshotFile] = []

# model is split, lets fetch all files based on the name pattern
match = re.match(SPLIT_MODEL_PATH_RE, self.model)
if match is None:
return files

path_part = match[1]
filename_base = match[2]
total_parts = int(match[3])

for i in range(1, total_parts + 1):
file_name = f"{filename_base}-{i:05d}-of-{total_parts:05d}.gguf"
url = f"{self.type}://{path_part}/{file_name}"
files.append(
SnapshotFile(
url=url,
header={},
hash=snapshot_hash,
type=SnapshotFileType.Model,
name=file_name,
should_show_progress=True,
required=True,
)
)

return files

def pull(self, _):
name, tag, _ = self.extract_model_identifiers()
_, _, all_files = self.model_store.get_cached_files(tag)
Expand All @@ -96,40 +125,21 @@ def pull(self, _):
self.model_store.new_snapshot(tag, snapshot_hash, files)
return

if not is_split_file_model(self.model):
files.append(
SnapshotFile(
url=f"{self.type}://{self.model}",
header={},
hash=snapshot_hash,
type=SnapshotFileType.Model,
name=name,
should_show_progress=True,
required=True,
)
)
if is_split_file_model(self.model):
files = self._assemble_split_file_list(snapshot_hash)
self.model_store.new_snapshot(tag, snapshot_hash, files)
return

# model is split, lets fetch all files based on the name pattern
match = re.match(SPLIT_MODEL_PATH_RE, self.model)
path_part = match[1]
filename_base = match[2]
total_parts = int(match[3])

for i in range(total_parts - 1):
i_off = i + 2
url = f"{self.type}://{path_part}/{filename_base}-{i_off:05d}-of-{total_parts:05d}.gguf"
files.append(
SnapshotFile(
url=url,
header={},
hash=snapshot_hash,
type=SnapshotFileType.Model,
name=name,
should_show_progress=True,
required=True,
)
files.append(
SnapshotFile(
url=f"{self.type}://{self.model}",
header={},
hash=snapshot_hash,
type=SnapshotFileType.Model,
name=name,
should_show_progress=True,
required=True,
)

)
self.model_store.new_snapshot(tag, snapshot_hash, files)
return
58 changes: 58 additions & 0 deletions test/unit/test_url.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
from dataclasses import dataclass, field

import pytest

from ramalama.model_store.snapshot_file import SnapshotFile
from ramalama.url import URL


@dataclass
class Input:
Model: str


@dataclass
class Expected:
URLList: list[str] = field(default_factory=lambda: [])
Names: list[str] = field(default_factory=lambda: [])


@pytest.mark.parametrize(
"input,expected",
[
(Input(""), Expected()),
(Input("file:///tmp/models/granite-3b-code-base.Q4_K_M.gguf"), Expected()),
(
Input(
"huggingface.co/unsloth/Qwen3-Coder-480B-A35B-Instruct-GGUF/resolve/main/Q3_K_M/Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00001-of-00005.gguf" # noqa: E501
),
Expected(
URLList=[
"https://huggingface.co/unsloth/Qwen3-Coder-480B-A35B-Instruct-GGUF/resolve/main/Q3_K_M/Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00001-of-00005.gguf", # noqa: E501
"https://huggingface.co/unsloth/Qwen3-Coder-480B-A35B-Instruct-GGUF/resolve/main/Q3_K_M/Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00002-of-00005.gguf", # noqa: E501
"https://huggingface.co/unsloth/Qwen3-Coder-480B-A35B-Instruct-GGUF/resolve/main/Q3_K_M/Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00003-of-00005.gguf", # noqa: E501
"https://huggingface.co/unsloth/Qwen3-Coder-480B-A35B-Instruct-GGUF/resolve/main/Q3_K_M/Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00004-of-00005.gguf", # noqa: E501
"https://huggingface.co/unsloth/Qwen3-Coder-480B-A35B-Instruct-GGUF/resolve/main/Q3_K_M/Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00005-of-00005.gguf", # noqa: E501
],
Names=[
"Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00001-of-00005.gguf",
"Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00002-of-00005.gguf",
"Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00003-of-00005.gguf",
"Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00004-of-00005.gguf",
"Qwen3-Coder-480B-A35B-Instruct-Q3_K_M-00005-of-00005.gguf",
],
),
),
],
)
def test__assemble_split_file_list(input: Input, expected: Expected):
# store path and scheme irrelevant here
model = URL(input.Model, "/store", "https")
files: list[SnapshotFile] = model._assemble_split_file_list("doesnotmatterhere")
file_count = len(files)
assert file_count == len(expected.Names)
assert file_count == len(expected.URLList)

for i in range(file_count):
assert files[i].url == expected.URLList[i]
assert files[i].name == expected.Names[i]
Loading