#!/usr/bin/env python3
"""Quantize Ornith-1.0-35B to 4-bit NF4 (recommended method for 2x RTX 5090)"""

from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoConfig
import torch
import os

def quantize_model():
    model_path = "/data/models/Ornith-1.0-35B"
    output_path = "/data/models/Ornith-1.0-35B-4bit-nf4"

    print(f"Quantizing model: {model_path}")
    print("Using 4-bit NF4 with double quantization + aggressive offloading...\n")

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )

    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        quantization_config=bnb_config,
        device_map="auto",
        max_memory={
            0: "26GiB",       # Good balance for 5090 (leaves headroom)
            1: "26GiB",
            "cpu": "150GiB",  # Heavy CPU offloading during quantization
        },
        low_cpu_mem_usage=True,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
    )

    print(f"\nSaving quantized model to: {output_path}")
    os.makedirs(output_path, exist_ok=True)
    model.save_pretrained(output_path)

    # Also save the config explicitly
    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
    config.save_pretrained(output_path)

    print(f"\n✅ Quantization complete!")
    print(f"   Model saved to: {output_path}")


if __name__ == "__main__":
    quantize_model()