Upload directory

main.py

@@ -0,0 +1,128 @@
+import torch
+from safetensors.torch import save_file, load_file
+import sys
+import json
+
+def create_gemma_safetensors(output_file="model.safetensors", dtype=torch.bfloat16, num_heads=64, num_layers=48, image_feature_dim=2048):
+    """
+    Creates a safetensors file with initialized Gemma model weights,
+    including vision components for a vision-text-to-text model.
+
+    Args:
+        output_file (str): The name of the output safetensors file.
+        dtype (torch.dtype): The data type of the model weights.
+        num_heads (int): The number of attention heads.
+        num_layers (int): The number of transformer layers.
+        image_feature_dim (int): The dimension of the image features.
+    """
+
+    tensors = {}
+    embedding_dim = 5376 // 4  # reducing embedding_dim to reduce model size
+    intermediate_size = 21504 // 4  # reducing intermediate_size to reduce model size
+    head_dim = embedding_dim // num_heads
+
+    # Text embedding
+    tensors["language_model.model.embed_tokens.weight"] = torch.randn(262208 // 4, embedding_dim, dtype=dtype) * 0.02
+
+    # Text layers
+    for layer_idx in range(num_layers):
+        layer_prefix = f"language_model.model.layers.{layer_idx}"
+        tensors[f"{layer_prefix}.input_layernorm.weight"] = torch.randn(embedding_dim, dtype=dtype)
+        tensors[f"{layer_prefix}.mlp.down_proj.weight"] = torch.randn(embedding_dim, intermediate_size, dtype=dtype) * (intermediate_size ** -0.5)
+        tensors[f"{layer_prefix}.mlp.gate_proj.weight"] = torch.randn(intermediate_size, embedding_dim, dtype=dtype) * (embedding_dim ** -0.5)
+        tensors[f"{layer_prefix}.mlp.up_proj.weight"] = torch.randn(intermediate_size, embedding_dim, dtype=dtype) * (embedding_dim ** -0.5)
+        tensors[f"{layer_prefix}.post_attention_layernorm.weight"] = torch.randn(embedding_dim, dtype=dtype)
+        tensors[f"{layer_prefix}.post_feedforward_layernorm.weight"] = torch.randn(embedding_dim, dtype=dtype)
+        tensors[f"{layer_prefix}.pre_feedforward_layernorm.weight"] = torch.randn(embedding_dim, dtype=dtype)
+        tensors[f"{layer_prefix}.self_attn.k_norm.weight"] = torch.randn(head_dim, dtype=dtype)
+        tensors[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn(head_dim * num_heads // 4, embedding_dim, dtype=dtype) * ((embedding_dim * head_dim // 4) ** -0.5)
+        tensors[f"{layer_prefix}.self_attn.o_proj.weight"] = torch.randn(embedding_dim, head_dim * num_heads // 4 * 2, dtype=dtype) * ((embedding_dim * head_dim // 4 * 2) ** -0.5)
+        tensors[f"{layer_prefix}.self_attn.q_norm.weight"] = torch.randn(head_dim, dtype=dtype)
+        tensors[f"{layer_prefix}.self_attn.q_proj.weight"] = torch.randn(head_dim * num_heads // 4 * 2, embedding_dim, dtype=dtype) * ((embedding_dim * head_dim // 4 * 2) ** -0.5)
+        tensors[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn(head_dim * num_heads // 4, embedding_dim, dtype=dtype) * ((embedding_dim * head_dim // 4) ** -0.5)
+
+    # Vision encoder (simplified example)
+    tensors["vision_encoder.projection.weight"] = torch.randn(image_feature_dim, embedding_dim, dtype=dtype) * (image_feature_dim ** -0.5)
+    tensors["vision_encoder.projection.bias"] = torch.zeros(embedding_dim, dtype=dtype)
+
+    total_params = 0
+    for tensor in tensors.values():
+        total_params += tensor.numel() * tensor.element_size()
+
+    total_params_gb = total_params / (1024 ** 3)
+
+    print(f"Estimated total parameter size: {total_params_gb:.2f} GB ({dtype}).")
+    proceed = input("Do you want to proceed and create the safetensors file? (y/n): ").lower()
+
+    if proceed == 'y':
+        save_file(tensors, output_file)
+        print(f"Gemma safetensors file created: {output_file}")
+    else:
+        print("Safetensors file creation cancelled.")
+        sys.exit()
+
+def analyze_model_file(model_file):
+    """
+    Analyzes a safetensors model file, extracts relevant information,
+    prints it to the console, and saves it to a config.json file.
+
+    Args:
+        model_file (str): The path to the safetensors model file.
+    """
+    try:
+        tensors = load_file(model_file, device="cpu")  # Load safetensors
+
+        config = {}
+
+        # Extract basic information (example, adjust based on your model)
+        config["architectures"] = ["Gemma3ForConditionalGeneration"]  # Changed architecture
+        config["model_type"] = "gemma3" # changed model_type
+        config["torch_dtype"] = str(tensors["language_model.model.embed_tokens.weight"].dtype)
+
+        # Extract dimensions (adjust based on your model's tensor names)
+        config["hidden_size"] = tensors["language_model.model.embed_tokens.weight"].shape[1]
+        config["num_hidden_layers"] = len(
+            [key for key in tensors if "language_model.model.layers." in key]
+        ) // 12
+        num_heads = tensors["language_model.model.layers.0.self_attn.q_proj.weight"].shape[0] // (config["hidden_size"] // 4)
+        config["num_attention_heads"] = num_heads
+        config["intermediate_size"] = tensors["language_model.model.layers.0.mlp.gate_proj.weight"].shape[0]
+        config["head_dim"] = config["hidden_size"] // config["num_attention_heads"]
+
+        # Add other relevant information as needed
+        # Example:
+        config["vocab_size"] = tensors["language_model.model.embed_tokens.weight"].shape[0]
+        config["image_feature_dim"] = tensors["vision_encoder.projection.weight"].shape[0]
+
+        # Print to console
+        print(json.dumps(config, indent=2))
+
+        # Save to config.json
+        with open("config.json", "w") as f:
+            json.dump(config, f, indent=2)
+
+        print("Model analysis completed and saved to config.json")
+
+    except FileNotFoundError:
+        print(f"Error: Model file '{model_file}' not found.")
+    except Exception as e:
+        print(f"An error occurred during model analysis: {e}")
+
+if __name__ == "__main__":
+    dtype_input = input("Enter the desired data type (bfloat16, float16, float32, etc.): ").lower()
+    if dtype_input == "bfloat16":
+        dtype = torch.bfloat16
+    elif dtype_input == "float16":
+        dtype = torch.float16
+    elif dtype_input == "float32":
+        dtype = torch.float32
+    elif dtype_input == "float8":
+        dtype = torch.float8_e4m3fn
+    else:
+        print("Invalid data type. Using bfloat16 as default.")
+        dtype = torch.bfloat16
+
+    create_gemma_safetensors(dtype=dtype)
+
+    # Analyze the created model file
+    analyze_model_file("model.safetensors")