Create create_safetensors.py

create_safetensors.py

@@ -0,0 +1,268 @@
+#!/usr/bin/env python3
+"""
+Create SafeTensors files for Mineral Nano 1 Vision (Mid-Range 4.1B)
+This creates randomly initialized weights ready for upload to Hugging Face
+"""
+
+import torch
+import json
+from safetensors.torch import save_file
+from pathlib import Path
+import math
+
+def initialize_weights(shape, init_type="normal", std=0.02):
+    """Initialize tensor with proper initialization"""
+    if init_type == "normal":
+        return torch.randn(shape) * std
+    elif init_type == "uniform":
+        limit = math.sqrt(3.0 * std)
+        return torch.rand(shape) * 2 * limit - limit
+    else:
+        return torch.zeros(shape)
+
+def create_mineral_nano_weights():
+    """Create all model weights for Mineral Nano 1 Vision"""
+    
+    print("Creating Mineral Nano 1 Vision SafeTensors (4.1B parameters)...")
+    print("This will take a few minutes and use ~16GB RAM temporarily...")
+    
+    state_dict = {}
+    
+    # Configuration
+    vocab_size = 64000
+    hidden_size = 2048
+    intermediate_size = 8192
+    num_layers = 28
+    num_heads = 32
+    num_kv_heads = 8
+    
+    vision_hidden = 1536
+    vision_intermediate = 6144
+    vision_layers = 24
+    vision_heads = 24
+    vision_patches = 1024  # 448/14 = 32, 32*32 = 1024
+    
+    # ============ LANGUAGE MODEL WEIGHTS ============
+    
+    print("\n[1/4] Creating language model embeddings...")
+    # Token embeddings
+    state_dict["model.embed_tokens.weight"] = initialize_weights(
+        (vocab_size, hidden_size), "normal", 0.02
+    )
+    
+    print("[2/4] Creating language model layers (28 layers)...")
+    # Transformer layers
+    for layer_idx in range(num_layers):
+        prefix = f"model.layers.{layer_idx}"
+        
+        # Self attention
+        state_dict[f"{prefix}.self_attn.q_proj.weight"] = initialize_weights(
+            (hidden_size, hidden_size), "normal", 0.02
+        )
+        state_dict[f"{prefix}.self_attn.k_proj.weight"] = initialize_weights(
+            (num_kv_heads * (hidden_size // num_heads), hidden_size), "normal", 0.02
+        )
+        state_dict[f"{prefix}.self_attn.v_proj.weight"] = initialize_weights(
+            (num_kv_heads * (hidden_size // num_heads), hidden_size), "normal", 0.02
+        )
+        state_dict[f"{prefix}.self_attn.o_proj.weight"] = initialize_weights(
+            (hidden_size, hidden_size), "normal", 0.02
+        )
+        
+        # MLP
+        state_dict[f"{prefix}.mlp.gate_proj.weight"] = initialize_weights(
+            (intermediate_size, hidden_size), "normal", 0.02
+        )
+        state_dict[f"{prefix}.mlp.up_proj.weight"] = initialize_weights(
+            (intermediate_size, hidden_size), "normal", 0.02
+        )
+        state_dict[f"{prefix}.mlp.down_proj.weight"] = initialize_weights(
+            (hidden_size, intermediate_size), "normal", 0.02
+        )
+        
+        # Layer norms
+        state_dict[f"{prefix}.input_layernorm.weight"] = torch.ones(hidden_size)
+        state_dict[f"{prefix}.post_attention_layernorm.weight"] = torch.ones(hidden_size)
+        
+        if (layer_idx + 1) % 5 == 0:
+            print(f"  ✓ Completed {layer_idx + 1}/{num_layers} layers")
+    
+    # Final layer norm and LM head
+    state_dict["model.norm.weight"] = torch.ones(hidden_size)
+    state_dict["lm_head.weight"] = initialize_weights(
+        (vocab_size, hidden_size), "normal", 0.02
+    )
+    
+    print("[3/4] Creating vision encoder (24 layers)...")
+    
+    # ============ VISION ENCODER WEIGHTS ============
+    
+    # Vision embeddings
+    state_dict["vision_model.embeddings.patch_embedding.weight"] = initialize_weights(
+        (vision_hidden, 3, 14, 14), "normal", 0.02
+    )
+    state_dict["vision_model.embeddings.class_embedding"] = initialize_weights(
+        (vision_hidden,), "normal", 0.02
+    )
+    state_dict["vision_model.embeddings.position_embedding.weight"] = initialize_weights(
+        (vision_patches + 1, vision_hidden), "normal", 0.02
+    )
+    
+    # Vision transformer layers
+    for layer_idx in range(vision_layers):
+        prefix = f"vision_model.encoder.layers.{layer_idx}"
+        
+        # Self attention
+        state_dict[f"{prefix}.self_attn.q_proj.weight"] = initialize_weights(
+            (vision_hidden, vision_hidden), "normal", 0.02
+        )
+        state_dict[f"{prefix}.self_attn.q_proj.bias"] = torch.zeros(vision_hidden)
+        
+        state_dict[f"{prefix}.self_attn.k_proj.weight"] = initialize_weights(
+            (vision_hidden, vision_hidden), "normal", 0.02
+        )
+        state_dict[f"{prefix}.self_attn.k_proj.bias"] = torch.zeros(vision_hidden)
+        
+        state_dict[f"{prefix}.self_attn.v_proj.weight"] = initialize_weights(
+            (vision_hidden, vision_hidden), "normal", 0.02
+        )
+        state_dict[f"{prefix}.self_attn.v_proj.bias"] = torch.zeros(vision_hidden)
+        
+        state_dict[f"{prefix}.self_attn.out_proj.weight"] = initialize_weights(
+            (vision_hidden, vision_hidden), "normal", 0.02
+        )
+        state_dict[f"{prefix}.self_attn.out_proj.bias"] = torch.zeros(vision_hidden)
+        
+        # MLP
+        state_dict[f"{prefix}.mlp.fc1.weight"] = initialize_weights(
+            (vision_intermediate, vision_hidden), "normal", 0.02
+        )
+        state_dict[f"{prefix}.mlp.fc1.bias"] = torch.zeros(vision_intermediate)
+        
+        state_dict[f"{prefix}.mlp.fc2.weight"] = initialize_weights(
+            (vision_hidden, vision_intermediate), "normal", 0.02
+        )
+        state_dict[f"{prefix}.mlp.fc2.bias"] = torch.zeros(vision_hidden)
+        
+        # Layer norms
+        state_dict[f"{prefix}.layer_norm1.weight"] = torch.ones(vision_hidden)
+        state_dict[f"{prefix}.layer_norm1.bias"] = torch.zeros(vision_hidden)
+        state_dict[f"{prefix}.layer_norm2.weight"] = torch.ones(vision_hidden)
+        state_dict[f"{prefix}.layer_norm2.bias"] = torch.zeros(vision_hidden)
+        
+        if (layer_idx + 1) % 5 == 0:
+            print(f"  ✓ Completed {layer_idx + 1}/{vision_layers} vision layers")
+    
+    # Vision post-layernorm
+    state_dict["vision_model.post_layernorm.weight"] = torch.ones(vision_hidden)
+    state_dict["vision_model.post_layernorm.bias"] = torch.zeros(vision_hidden)
+    
+    print("[4/4] Creating vision-language projection...")
+    
+    # ============ MULTI-MODAL PROJECTOR ============
+    
+    # Project vision features to language model dimension
+    state_dict["multi_modal_projector.linear_1.weight"] = initialize_weights(
+        (hidden_size, vision_hidden), "normal", 0.02
+    )
+    state_dict["multi_modal_projector.linear_1.bias"] = torch.zeros(hidden_size)
+    
+    state_dict["multi_modal_projector.linear_2.weight"] = initialize_weights(
+        (hidden_size, hidden_size), "normal", 0.02
+    )
+    state_dict["multi_modal_projector.linear_2.bias"] = torch.zeros(hidden_size)
+    
+    return state_dict
+
+def save_sharded_safetensors(state_dict, output_dir="./mineral-nano-weights"):
+    """Save model in sharded format for Hugging Face"""
+    
+    output_path = Path(output_dir)
+    output_path.mkdir(exist_ok=True)
+    
+    print(f"\nSaving SafeTensors files to {output_dir}/")
+    
+    # Calculate sizes and split into shards
+    max_shard_size = 4.5 * 1024 * 1024 * 1024  # 4.5 GB per shard
+    
+    shards = []
+    current_shard = {}
+    current_size = 0
+    shard_idx = 1
+    
+    for key, tensor in state_dict.items():
+        tensor_size = tensor.numel() * tensor.element_size()
+        
+        if current_size + tensor_size > max_shard_size and current_shard:
+            # Save current shard
+            shard_file = f"model-{shard_idx:05d}-of-00002.safetensors"
+            print(f"  Saving {shard_file} ({current_size / 1e9:.2f} GB)...")
+            save_file(current_shard, output_path / shard_file)
+            shards.append((shard_file, list(current_shard.keys())))
+            
+            # Start new shard
+            current_shard = {}
+            current_size = 0
+            shard_idx += 1
+        
+        current_shard[key] = tensor
+        current_size += tensor_size
+    
+    # Save final shard
+    if current_shard:
+        shard_file = f"model-{shard_idx:05d}-of-00002.safetensors"
+        print(f"  Saving {shard_file} ({current_size / 1e9:.2f} GB)...")
+        save_file(current_shard, output_path / shard_file)
+        shards.append((shard_file, list(current_shard.keys())))
+    
+    # Create index file
+    index = {
+        "metadata": {
+            "total_size": sum(t.numel() * t.element_size() for t in state_dict.values())
+        },
+        "weight_map": {}
+    }
+    
+    for shard_file, keys in shards:
+        for key in keys:
+            index["weight_map"][key] = shard_file
+    
+    index_path = output_path / "model.safetensors.index.json"
+    with open(index_path, "w") as f:
+        json.dump(index, f, indent=2)
+    
+    print(f"  ✓ Saved index file: model.safetensors.index.json")
+    
+    # Calculate total parameters
+    total_params = sum(t.numel() for t in state_dict.values())
+    total_size = sum(t.numel() * t.element_size() for t in state_dict.values())
+    
+    print(f"\n{'='*60}")
+    print(f"✓ SUCCESS! Model weights created:")
+    print(f"  Total Parameters: {total_params:,} ({total_params/1e9:.2f}B)")
+    print(f"  Total Size: {total_size/1e9:.2f} GB")
+    print(f"  Number of Shards: {len(shards)}")
+    print(f"  Output Directory: {output_path.absolute()}")
+    print(f"{'='*60}")
+    print(f"\nNext steps:")
+    print(f"1. Upload these files to Hugging Face:")
+    print(f"   - model-00001-of-00002.safetensors")
+    print(f"   - model-00002-of-00002.safetensors")
+    print(f"   - model.safetensors.index.json")
+    print(f"2. Your model will be ready to load!")
+    print(f"\nNote: These are randomly initialized weights.")
+    print(f"For a working model, you need to train or fine-tune.")
+
+if __name__ == "__main__":
+    print("="*60)
+    print("Mineral Nano 1 Vision - SafeTensors Generator")
+    print("Mid-Range 4.1B Parameter Model")
+    print("="*60)
+    
+    # Create weights
+    state_dict = create_mineral_nano_weights()
+    
+    # Save to disk
+    save_sharded_safetensors(state_dict)
+    
+    print("\n✓ All done! Your SafeTensors files are ready.")