import os
import torch
from models import FromZero
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)

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

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


if __name__ == "__main__":
    main()