#!/usr/bin/env python3
"""Convert torch.save quantized shards to safetensors format."""

import argparse
import gc
import torch
from pathlib import Path
from safetensors.torch import save_file
import glob


def convert_shards(model_path):
    output_path = Path(model_path)
    shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
    
    print(f"Converting {len(shards)} shards to safetensors format...\n")
    
    for i, shard_path in enumerate(shards):
        print(f"Converting shard {i+1}/{len(shards)}: {Path(shard_path).name}")
        
        # Load torch.save format
        ckpt = torch.load(shard_path, map_location="cpu", weights_only=False)
        
        # Separate tensors from QuantState objects
        tensors = {}
        for key, value in ckpt.items():
            if isinstance(value, torch.Tensor):
                tensors[key] = value
            else:
                # Save QuantState separately if needed
                print(f"   Skipping non-tensor: {key} ({type(value)})")
        
        # Save as safetensors
        new_path = output_path / Path(shard_path).name
        save_file(tensors, new_path)
        print(f"   ✓ Saved {len(tensors)} tensors")
        
        # Cleanup
        del ckpt, tensors
        gc.collect()
    
    print(f"\n✅ Converted {len(shards)} shards to safetensors format")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-path", type=str, default="/data/models/Ornith-1.0-35B-nf4")
    args = parser.parse_args()
    convert_shards(args.model_path)