from transformers import AutoTokenizer
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
import torch

# 1. Configuración
model_id = "prometheus-eval/prometheus-7b-v2.0"
output_dir = "prometheus-7b-v2.0-GPTQ-4bit"

# Configuración de 4 bits
quantize_config = BaseQuantizeConfig(
    bits=4,          # Cuantización a 4 bits
    group_size=128,  # Recomendado para equilibrio calidad/velocidad
    desc_act=False,  # Mejora la velocidad de inferencia
)

# 2. Cargar Tokenizer y preparar datos de calibración
# GPTQ necesita unos pocos ejemplos para ajustar los pesos
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
examples = [
    tokenizer("Prometheus is a specialized language model for evaluating other AI models."),
    tokenizer("Quantization helps to run large models on consumer hardware.")
]

# 3. Cargar el modelo original (FP16)
model = AutoGPTQForCausalLM.from_pretrained(
    model_id,
    quantize_config,
    device_map="auto",
    torch_dtype=torch.float16
)

# 4. Ejecutar la cuantización
model.quantize(examples)

# 5. Guardar el modelo resultante en formato .safetensors
model.save_quantized(output_dir, use_safetensors=True)
tokenizer.save_pretrained(output_dir)

print(f"Modelo guardado exitosamente en: {output_dir}")