import os
import torch
from transformers import LlamaConfig, LlamaForCausalLM
import accelerate

def initialize():
    # Load the 350M config we just created
    model_dir = "e:/agent/agent_ai/sail/sail_hf_model" if os.name == 'nt' else "/mnt/e/agent/agent_ai/sail/sail_hf_model"
    
    print(f"Loading 350M configuration from {model_dir}/config.json ...")
    config = LlamaConfig.from_pretrained(model_dir)
    
    print("Initializing weights from scratch for 350M parameters (this will utilize ~1.4GB)...")
    # We initialize the model instantaneously using accelerate
    with accelerate.init_empty_weights():
        model = LlamaForCausalLM(config)
    
    # Materialize weights (creates random normal distribution internally during instantiation)
    model.to_empty(device="cpu")
    
    # Ensure standard normal initialization for parameters
    for param in model.parameters():
        torch.nn.init.normal_(param, mean=0.0, std=0.02)
        
    print(f"Total parameters: {model.num_parameters():,}")
    
    # Save using safe_serialization (safetensors format), which is standard for Unsloth / Llama Factory
    print("Saving to safetensors format...")
    model.save_pretrained(model_dir, safe_serialization=True)
    
    print("DONE! Your 350M foundational blank model is ready in safetensors format.")
    print("You can now load this directory directly into Unsloth or Llama-Factory for pre-training.")

if __name__ == "__main__":
    initialize()