vrr/kvcache-ai-ktransformers
2
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2025-09-06 04:30:03 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

get dirname if gguf_path is a file

Browse source
This commit is contained in:
Shuaiyi 2025-02-14 10:22:56 +00:00
parent 1548c99234
commit 22280bf17f
1 changed files with 6 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
ktransformers/util/custom_gguf.py
View file

@ -166,6 +166,8 @@ class GGUFLoader:
# Check dir exist
if not os.path.exists(gguf_path):
raise FileNotFoundError(f"GGUF dir not found: {gguf_path}")
if os.path.isfile(gguf_path):
gguf_path = os.path.dirname(gguf_path)
self.tensor_info = {}
self.gguf_path = gguf_path
@ -175,14 +177,18 @@ class GGUFLoader:
self.tensor_device_map = {}
# Walk through all the .gguf files in the directory
found_gguf = False
for root, dirs, files in os.walk(gguf_path):
for file in files:
if file.endswith(".gguf"):
found_gguf = True
file_name = os.path.join(root, file)
with open(file_name, "rb") as f:
self.load_gguf(f)
if file_name not in self.file_data_map:
self.file_data_map[file_name] = np.memmap(file_name, mode = 'r')
if not found_gguf:
raise FileNotFoundError(f"Cannot find any .gguf files in: {gguf_path}")
def load_gguf(self, f):
f.seek(0)