rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Universal ziproot for lazy_loader

Browse Source
This commit is contained in:
Henk
2023-07-07 13:37:32 +02:00
parent f0d161b9c6
commit 76d21bb142
1 changed files with 7 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
modeling/lazy_loader.py
View File

@@ -172,14 +172,11 @@ class TorchLazyTensor(LazyTensor):
CheckpointChunkCache.file_name = filename CheckpointChunkCache.file_name = filename
CheckpointChunkCache.key = self.key CheckpointChunkCache.key = self.key
try: ziproot = checkpoint.namelist()[0].split("/")[0]
CheckpointChunkCache.handle = checkpoint.open( CheckpointChunkCache.handle = checkpoint.open(f"{ziproot}/data/{self.key}", "r")
f"archive/data/{self.key}", "r"
)
except KeyError:
CheckpointChunkCache.handle = checkpoint.open(
f"{filename}/data/{self.key}", "r"
)
else: else:
# Cache hit. Hip hip hooray! :^) # Cache hit. Hip hip hooray! :^)
# print(".", end="", flush=True) # print(".", end="", flush=True)
@@ -520,8 +517,8 @@ def use_lazy_load(
torch.load = torch_load torch.load = torch_load
if HAS_SAFETENSORS: #if HAS_SAFETENSORS:
patch_safetensors(callback) #patch_safetensors(callback)
if dematerialized_modules: if dematerialized_modules:
# Most devices can just use Accelerate's implementation, but the Transformers on # Most devices can just use Accelerate's implementation, but the Transformers on