fix: add index file creation for sharded models

train.py

@@ -33,6 +33,42 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
+    # Create index file for sharded model
+    import glob as glob_mod
+    import json as json_mod
+    from safetensors import safe_open
+    
+    safetensor_files = glob_mod.glob(f"{config['base_model']}/*.safetensors")
+    shards = sorted([Path(f).name for f in safetensor_files if "of-" in Path(f).name])
+    index_file = Path(config["base_model"]) / "model.safetensors.index.json"
+    
+    if not index_file.exists() and shards:
+        print(f"\n[INFO] Creating index file for {len(shards)} shards...")
+        weight_map = {}
+        for shard_name in shards:
+            shard_path = Path(config["base_model"]) / shard_name
+            with safe_open(str(shard_path), framework="pt", device="cpu") as f:
+                for key in f.keys():
+                    weight_map[key] = shard_name
+        
+        index = {
+            "metadata": {"total_size": sum((Path(config["base_model"]) / s).stat().st_size for s in shards)},
+            "weight_map": weight_map
+        }
+        with open(index_file, 'w') as f:
+            json_mod.dump(index, f)
+        print(f"✓ Created index ({len(weight_map)} weights)")
+    
+    # Remove quantization_config to prevent re-quantization
+    config_json_path = Path(config["base_model"]) / "config.json"
+    if config_json_path.exists():
+        with open(config_json_path, 'r') as f:
+            config_data = json_mod.load(f)
+        if 'quantization_config' in config_data:
+            del config_data['quantization_config']
+            with open(config_json_path, 'w') as f:
+                json_mod.dump(config_data, f)
+    
     print(f"\n[INFO] Loading pre-quantized BnB 4-bit model...")
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],