#!/usr/bin/env python3
"""
Workaround script to convert Comma v0.1-2T to GGUF format.
This uses transformers to load the model, then saves it in a format
that llama.cpp can better handle.
"""

import sys
import os
from pathlib import Path

try:
    from transformers import AutoModelForCausalLM, AutoTokenizer
    import torch
except ImportError:
    print("ERROR: Please install transformers: pip install transformers torch")
    sys.exit(1)

model_path = Path("comma-v0.1-2t")
output_path = Path("comma-v0.1-2t-converted")

print(f"Loading model from {model_path}...")
print("This may take several minutes...")

# Load model and tokenizer
try:
    model = AutoModelForCausalLM.from_pretrained(
        str(model_path),
        torch_dtype=torch.float16,
        device_map="cpu",  # Keep on CPU for conversion
        low_cpu_mem_usage=True
    )
    tokenizer = AutoTokenizer.from_pretrained(str(model_path))

    print(f"Model loaded successfully!")
    print(f"Vocabulary size: {len(tokenizer)}")
    print(f"Model parameters: {sum(p.numel() for p in model.parameters()) / 1e9:.2f}B")

    # Save in a clean format
    output_path.mkdir(exist_ok=True)
    print(f"\nSaving converted model to {output_path}...")

    model.save_pretrained(str(output_path), safe_serialization=True)
    tokenizer.save_pretrained(str(output_path))

    print("\nConversion complete!")
    print(f"Now try: python llama.cpp/convert_hf_to_gguf.py {output_path} --outfile comma-v0.1-2t.gguf --outtype q4_K_M")

except Exception as e:
    print(f"ERROR: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)