Rename model to model-1-of-10.safetensors

model-1-of-10.safetensors

@@ -0,0 +1,60 @@
+from safetensors.torch import load_file, save_file
+import torch
+from typing import List, Dict
+import logging
+from tqdm import tqdm
+import os
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+def load_charm_model(model_parts: List[str]) -> Dict[str, torch.Tensor]:
+    """
+    Load and merge multiple .safetensors model files.
+
+    Args:
+        model_parts (list): List of model part file paths (e.g., ["model-1-of-10.safetensors", ...])
+
+    Returns:
+        dict: Merged model state dictionary
+
+    Raises:
+        FileNotFoundError: If any model part file is missing.
+        RuntimeError: If there is an issue loading or merging the model parts.
+    """
+    merged_state_dict = {}
+
+    # Check if all model parts exist
+    for part in model_parts:
+        if not os.path.exists(part):
+            raise FileNotFoundError(f"Model part not found: {part}")
+
+    # Load and merge model parts
+    try:
+        logger.info("Loading and merging model parts...")
+        for part in tqdm(model_parts, desc="Loading model parts"):
+            state_dict = load_file(part)
+            merged_state_dict.update(state_dict)  # Merge parameters
+            logger.debug(f"Loaded part: {part}")
+
+        logger.info("Model parts loaded and merged successfully.")
+        return merged_state_dict
+    except Exception as e:
+        logger.error(f"Error loading or merging model parts: {e}")
+        raise RuntimeError("Failed to load or merge model parts.")
+
+# Example usage
+if __name__ == "__main__":
+    try:
+        # List of model part files
+        model_files = [f"model-{i}-of-10.safetensors" for i in range(1, 11)]
+
+        # Load and merge the model
+        charm_model = load_charm_model(model_files)
+
+        # Save the merged model as a .safetensors file
+        save_file(charm_model, "merged_model.safetensors")
+        logger.info("Merged model saved as 'merged_model.safetensors'.")
+    except Exception as e:
+        logger.error(f"An error occurred: {e}")