39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
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}") |