import os
import torch
from models import FromZero, PreTrained
from safetensors.torch import save_model
from utils import cargar_etiquetas

def main():
    # Crear la carpeta model_test si no existe
    os.makedirs("model_test", exist_ok=True)
    _,num_classes,_ = cargar_etiquetas()
    # Crear instancias de los modelos
    from_zero_model = FromZero(num_classes=num_classes)
    pretrained_model = PreTrained(num_classes=num_classes)

    # Guardar los modelos
    save_model(from_zero_model, "model_test/from_zero_model.safetensor")
    save_model(pretrained_model, "model_test/pretrained_model.safetensor")

    print("Los modelos han sido creados y guardados en la carpeta 'model_test'")


if __name__ == "__main__":
    main()