unifiedefficientloader 0.2.0__tar.gz → 0.2.1__tar.gz

{unifiedefficientloader-0.2.0 → unifiedefficientloader-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unifiedefficientloader
-Version: 0.2.0
+Version: 0.2.1
 Summary: A unified interface for memory efficient per tensor loading of safetensors files as raw bytes from offset, handling CPU/GPU pinned transfers, and converting between tensors and dicts.
 Author: silveroxides
 License: MIT

{unifiedefficientloader-0.2.0 → unifiedefficientloader-0.2.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "unifiedefficientloader"
-version = "0.2.0"
+version = "0.2.1"
 description = "A unified interface for memory efficient per tensor loading of safetensors files as raw bytes from offset, handling CPU/GPU pinned transfers, and converting between tensors and dicts."
 readme = "README.md"
 authors = [

{unifiedefficientloader-0.2.0 → unifiedefficientloader-0.2.1}/unifiedefficientloader/memory_efficient_loader.py

@@ -244,6 +244,30 @@ class UnifiedSafetensorsLoader:
         else:
             raise ValueError(f"Unsupported float8 type: {dtype_str}")
 
+
+    def load_all(self):
+        """Load all tensors as a dictionary.
+
+        In low-memory mode: uses async_stream for parallel multi-threaded I/O.
+        In standard mode: returns the preloaded tensor cache.
+
+        Returns:
+            Dict[str, torch.Tensor]: All tensors keyed by name.
+        """
+        if not self.low_memory:
+            return dict(self._tensors)
+
+        sd = {}
+        for batch in self.async_stream(
+            keys=self._all_keys,
+            batch_size=16,
+            prefetch_batches=2,
+            pin_memory=False,
+        ):
+            for key, tensor in batch:
+                sd[key] = tensor
+        return sd
+
     def async_stream(self, keys: list, batch_size: int = 1, prefetch_batches: int = 2, pin_memory: bool = False):
         """Asynchronously stream tensors from disk.
         

{unifiedefficientloader-0.2.0 → unifiedefficientloader-0.2.1}/unifiedefficientloader.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unifiedefficientloader
-Version: 0.2.0
+Version: 0.2.1
 Summary: A unified interface for memory efficient per tensor loading of safetensors files as raw bytes from offset, handling CPU/GPU pinned transfers, and converting between tensors and dicts.
 Author: silveroxides
 License: MIT