File size: 2,930 Bytes
4d061f7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import logging
from huggingface_hub import HfApi, create_repo, upload_folder
from transformers import AutoConfig, AutoModelForCausalLM, PreTrainedModel
from transformers.models.llama.modeling_llama import LlamaConfig, LlamaForCausalLM
from torch import nn
# Importar BitLinear desde el m贸dulo de cuantizaci贸n existente
from src.training.utils.quantization import BitLinear
from src.training.utils.linear_to_bitlinear import replace_linears_in_hf
# Cuantizaci贸n y conversi贸n a BitNet
class BitNetConfig(LlamaConfig):
model_type = "bitnet"
def __init__(self, **kwargs):
super().__init__(**kwargs)
class BitNetModel(PreTrainedModel):
config_class = BitNetConfig
def __init__(self, config):
super().__init__(config)
self.model = LlamaForCausalLM(config)
def forward(self, *args, **kwargs):
return self.model(*args, **kwargs)
# Registrar BitNet como un nuevo tipo de modelo en Hugging Face Transformers
AutoConfig.register("bitnet", BitNetConfig)
AutoModelForCausalLM.register(BitNetConfig, BitNetModel)
def save_convert_and_push_model(trainer, output_path, huggingface_id, new_model_name, hf_token, convert_to_bitnet=False, do_push=True):
"""
Convierte el modelo a BitNet si se solicita, guarda el modelo y opcionalmente lo sube a Hugging Face Hub.
Args:
trainer: Trainer instance con el modelo ya entrenado.
output_path: Ruta para guardar el modelo localmente.
huggingface_id: ID del usuario u organizaci贸n en Hugging Face Hub.
new_model_name: Nombre del nuevo modelo que se subir谩 al Hub.
hf_token: Token de acceso para Hugging Face Hub.
convert_to_bitnet (bool): Flag para decidir si se convierte a BitNet.
do_push (bool): Flag para decidir si se hace push al Hugging Face Hub.
Returns:
None
"""
logger = logging.getLogger(__name__)
# Convertir a BitNet si se especifica
if convert_to_bitnet:
logger.info("馃攧 Converting the model to BitNet architecture.")
replace_linears_in_hf(trainer.model)
output_dir = f"{output_path}/bitnet_model"
else:
output_dir = f"{output_path}/final_model"
# Guardar el modelo localmente
logger.info(f"馃捑 Saving the model locally to {output_dir}.")
trainer.model.save_pretrained(output_dir)
trainer.tokenizer.save_pretrained(output_dir)
# Subir el modelo al Hugging Face Hub si se solicita
if do_push:
logger.info("鈽侊笍 Uploading the model and tokenizer to Hugging Face Hub.")
api = HfApi()
create_repo(
repo_id=f"{huggingface_id}/{new_model_name}",
repo_type="model",
exist_ok=True,
token=hf_token,
)
upload_folder(
folder_path=output_dir,
repo_type="model",
repo_id=f"{huggingface_id}/{new_model_name}",
token=hf_token,
)
|