Delete main.py

main.py

@@ -1,219 +0,0 @@
-import torch
-from safetensors.torch import save_file, load_file
-import sys
-import json
-
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-def create_bitnet_safetensors(output_file="model_1.safetensors", 
-                             dtype=torch.bfloat16, 
-                             num_heads=4, 
-                             num_layers=4, 
-                             image_feature_dim=256, 
-                             device=device,
-                             quantize=True):
-    """
-    Creates a safetensors file with initialized BitNet model weights.
-    """
-    tensors = {}
-
-    # These values reduce model size; adjust as needed
-    embedding_dim = 5376 // 4
-    intermediate_size = 21504 // 4
-    head_dim = embedding_dim // num_heads
-
-    # Text Encoder Components
-    # Text embedding
-    tensors["language_model.model.embed_tokens.weight"] = torch.randn(262208 // 4, embedding_dim, dtype=dtype, device=device) * 0.02
-
-    # Text layers
-    for layer_idx in range(num_layers):
-        layer_prefix = f"language_model.model.layers.{layer_idx}"
-
-        # Layer normalization
-        tensors[f"{layer_prefix}.pre_attn_layernorm.weight"] = torch.ones(embedding_dim, dtype=dtype, device=device)
-        tensors[f"{layer_prefix}.post_mlp_layernorm.weight"] = torch.ones(embedding_dim, dtype=dtype, device=device)
-
-        # Self-attention projections (Q, K, V)
-        tensors[f"{layer_prefix}.self_attn.q_proj.weight"] = torch.randn(embedding_dim, embedding_dim, dtype=dtype, device=device) * (embedding_dim ** -0.5)
-        tensors[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn(embedding_dim, embedding_dim, dtype=dtype, device=device) * (embedding_dim ** -0.5)
-        tensors[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn(embedding_dim, embedding_dim, dtype=dtype, device=device) * (embedding_dim ** -0.5)
-
-        # Self-attention output projection
-        tensors[f"{layer_prefix}.self_attn.out_proj.weight"] = torch.randn(embedding_dim, embedding_dim, dtype=dtype, device=device) * (embedding_dim ** -0.5)
-
-        # MLP sub-block: first linear layer
-        tensors[f"{layer_prefix}.mlp.fc1.weight"] = torch.randn(intermediate_size, embedding_dim, dtype=dtype, device=device) * (embedding_dim ** -0.5)
-        tensors[f"{layer_prefix}.mlp.fc1.bias"] = torch.zeros(intermediate_size, dtype=dtype, device=device)
-
-        # MLP sub-block: second linear layer
-        tensors[f"{layer_prefix}.mlp.fc2.weight"] = torch.randn(embedding_dim, intermediate_size, dtype=dtype, device=device) * (intermediate_size ** -0.5)
-        tensors[f"{layer_prefix}.mlp.fc2.bias"] = torch.zeros(embedding_dim, dtype=dtype, device=device)
-
-    # Vision Encoder Components
-    patch_size = 16
-    image_size = 224
-    num_patches = (image_size // patch_size) ** 2
-    vision_embedding_dim = embedding_dim
-
-    # Patch Embedding
-    tensors["vision_encoder.patch_embedding.weight"] = torch.randn(vision_embedding_dim, 3, patch_size, patch_size, dtype=dtype, device=device) * 0.02
-    tensors["vision_encoder.patch_embedding.bias"] = torch.zeros(vision_embedding_dim, dtype=dtype, device=device)
-
-    # Positional Embeddings
-    tensors["vision_encoder.position_embeddings"] = torch.randn(1, num_patches + 1, vision_embedding_dim, dtype=dtype, device=device) * 0.02
-
-    # Class Token
-    tensors["vision_encoder.cls_token"] = torch.randn(1, 1, vision_embedding_dim, dtype=dtype, device=device) * 0.02
-
-    # Vision Transformer Layers
-    for layer_idx in range(num_layers):
-        layer_prefix = f"vision_encoder.transformer.layers.{layer_idx}"
-
-        # Pre-Attention Layer Normalization
-        tensors[f"{layer_prefix}.pre_attn_layernorm.weight"] = torch.ones(vision_embedding_dim, dtype=dtype, device=device)
-
-        # Self-Attention projections (Q, K, V)
-        tensors[f"{layer_prefix}.self_attn.q_proj.weight"] = torch.randn(vision_embedding_dim, vision_embedding_dim, dtype=dtype, device=device) * (vision_embedding_dim ** -0.5)
-        tensors[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn(vision_embedding_dim, vision_embedding_dim, dtype=dtype, device=device) * (vision_embedding_dim ** -0.5)
-        tensors[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn(vision_embedding_dim, vision_embedding_dim, dtype=dtype, device=device) * (vision_embedding_dim ** -0.5)
-
-        # Self-Attention output projection
-        tensors[f"{layer_prefix}.self_attn.out_proj.weight"] = torch.randn(vision_embedding_dim, vision_embedding_dim, dtype=dtype, device=device) * (vision_embedding_dim ** -0.5)
-
-        # MLP sub-block: first linear layer
-        tensors[f"{layer_prefix}.mlp.fc1.weight"] = torch.randn(vision_embedding_dim, intermediate_size, dtype=dtype, device=device) * (vision_embedding_dim ** -0.5)
-        tensors[f"{layer_prefix}.mlp.fc1.bias"] = torch.zeros(intermediate_size, dtype=dtype, device=device)
-
-        # MLP sub-block: second linear layer
-        tensors[f"{layer_prefix}.mlp.fc2.weight"] = torch.randn(intermediate_size, vision_embedding_dim, dtype=dtype, device=device) * (vision_embedding_dim ** -0.5)
-        tensors[f"{layer_prefix}.mlp.fc2.bias"] = torch.zeros(vision_embedding_dim, dtype=dtype, device=device)
-
-        # Post-MLP Layer Normalization
-        tensors[f"{layer_prefix}.post_mlp_layernorm.weight"] = torch.ones(vision_embedding_dim, dtype=dtype, device=device)
-
-    total_params = sum(tensor.numel() * tensor.element_size() for tensor in tensors.values())
-    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':
-        if quantize:
-            # Apply 1.58-bit quantization
-            quantized_tensors = {}
-            for key, tensor in tensors.items():
-                # Normalize the tensor
-                max_val = torch.max(torch.abs(tensor))
-                normalized = tensor / max_val
-                
-                # Quantize to {-1, 0, 1}
-                quantized = torch.where(normalized > 0.5, torch.ones_like(normalized), 
-                                      torch.where(normalized < -0.5, -torch.ones_like(normalized), torch.zeros_like(normalized)))
-                
-                quantized_tensors[key] = quantized
-                print(f"Quantized {key} to 1.58-bit precision.")
-            
-            tensors = quantized_tensors
-
-        tensors = {k: v.cpu() for k, v in tensors.items()}
-        save_file(tensors, output_file)
-        print(f"BitNet safetensors file created: {output_file}")
-    else:
-        print("Safetensors file creation cancelled.")
-        sys.exit()
-
-def quantize_model(model_file, output_file, device=device):
-    """
-    Quantizes the given model to 1.58-bit precision.
-    """
-    try:
-        # Load the model
-        tensors = load_file(model_file, device=device)
-        
-        # Quantize each tensor
-        quantized_tensors = {}
-        for key, tensor in tensors.items():
-            # Normalize the tensor
-            max_val = torch.max(torch.abs(tensor))
-            normalized = tensor / max_val
-            
-            # Quantize to {-1, 0, 1}
-            quantized = torch.where(normalized > 0.5, torch.ones_like(normalized), 
-                                  torch.where(normalized < -0.5, -torch.ones_like(normalized), torch.zeros_like(normalized)))
-            
-            quantized_tensors[key] = quantized
-            print(f"Quantized {key} to 1.58-bit precision.")
-        
-        # Save the quantized model
-        save_file(quantized_tensors, output_file)
-        print(f"Quantized model saved to {output_file}")
-        
-    except FileNotFoundError:
-        print(f"Error: Model file '{model_file}' not found.")
-    except Exception as e:
-        print(f"An error occurred during quantization: {e}")
-
-def analyze_model_file(model_file):
-    """
-    Analyzes a safetensors model file and prints/saves the configuration.
-    """
-    try:
-        tensors = load_file(model_file, device="cpu")
-        config = {}
-
-        # Basic model architecture details
-        config["architectures"] = ["TestLakeForConditionalGeneration"]
-        config["model_type"] = "testlake"
-        config["torch_dtype"] = str(tensors["language_model.model.embed_tokens.weight"].dtype)
-
-        # Extract dimensions for language model
-        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.fc1.weight"].shape[1]
-        config["head_dim"] = config["hidden_size"] // config["num_attention_heads"]
-
-        # Extract vocabulary and vision details
-        config["vocab_size"] = tensors["language_model.model.embed_tokens.weight"].shape[0]
-        config["image_feature_dim"] = tensors["vision_encoder.patch_embedding.weight"].shape[0]
-        config["num_patches"] = ((224 // 16) ** 2)
-        config["vision_num_layers"] = len([key for key in tensors if "vision_encoder.transformer.layers." in key]) // 6
-
-        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__":
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    
-    # Set the desired data type for training (BFloat16)
-    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
-    else:
-        print("Invalid data type. Using bfloat16 as default.")
-        dtype = torch.bfloat16
-
-    # Create the initial model with BFloat16 precision and quantize it
-    create_bitnet_safetensors(dtype=dtype, device=device, quantize=True)
-
-    # Analyze the created model file
-    analyze_model_file("model_1.safetensors")