import json
from safetensors import safe_open

def generate_safetensors_index(model_path="."):
    """Generate model.safetensors.index.json from existing safetensors files"""
    
    # Load the existing bin index as reference
    with open(f"pytorch_model.bin.index.json", "r") as f:
        bin_index = json.load(f)
    
    # Initialize the safetensors index structure
    safetensors_index = {
        "metadata": bin_index.get("metadata", {}),
        "weight_map": {}
    }
    
    # Map each safetensors file and get its tensor names
    safetensors_files = [
        "pytorch_model-00001-of-00004.safetensors",
        "pytorch_model-00002-of-00004.safetensors", 
        "pytorch_model-00003-of-00004.safetensors",
        "pytorch_model-00004-of-00004.safetensors"
    ]
    
    for safetensor_file in safetensors_files:
        try:
            with safe_open(f"{safetensor_file}", framework="pt") as f:
                for tensor_name in f.keys():
                    safetensors_index["weight_map"][tensor_name] = safetensor_file
            print(f"✓ Processed {safetensor_file}")
        except Exception as e:
            print(f"✗ Error processing {safetensor_file}: {e}")
    
    # Save the index file
    with open(f"model.safetensors.index.json", "w") as f:
        json.dump(safetensors_index, f, indent=2)
    
    print(f"✓ Generated model.safetensors.index.json with {len(safetensors_index['weight_map'])} tensors")
    return safetensors_index

# Run the function
if __name__ == "__main__":
    # Change this path to your model directory if needed
    generate_safetensors_index("./Finance-Llama-8B")