File size: 2,485 Bytes
fc4232d 631d039 fc4232d 631d039 fc4232d 631d039 fc4232d afed862 631d039 fc4232d afed862 fc4232d afed862 fc4232d |
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 os
import base64
import shutil
from huggingface_hub import HfApi, upload_file
def main():
# Obtiene el nombre del modelo de la variable de entorno
model_name = os.getenv('MODELNAME')
if model_name is None:
print("MODELNAME no está definida en las variables de entorno.")
return
# Ruta a los archivos
onnx_file = f"{model_name}.onnx"
config_file = "config.json"
onnx_json_file = f"{model_name}.onnx.json"
# Verifica si el archivo .onnx existe
if os.path.isfile(onnx_file):
print(f"{model_name}.onnx encontrado.")
else:
print(f"{model_name}.onnx no se encuentra.")
return
# Renombra config.json a ${MODELNAME}.onnx.json si existe
if os.path.isfile(config_file):
shutil.move(config_file, onnx_json_file)
print(f"Renombrado config.json a {onnx_json_file}")
else:
print("config.json no se encuentra.")
return
# Obtén el token desde la variable de entorno llamada 'TOKEN'
token_base64 = os.getenv('TOKEN')
if token_base64 is None:
print("TOKEN no está definida en las variables de entorno.")
return
# Decodifica el token de base64
token = base64.b64decode(token_base64).decode('utf-8')
# Configuración para la API de Hugging Face
api = HfApi()
# Obtiene el nombre de usuario
user_info = api.whoami(token=token)
username = user_info['name']
# Crea el repositorio en Hugging Face si no existe
repo_id = f"{username}/{model_name}"
try:
api.create_repo(repo_id=repo_id, token=token)
print(f"Repositorio {repo_id} creado en Hugging Face.")
except Exception as e:
print(f"El repositorio {repo_id} ya existe o se produjo un error: {e}")
# Subir archivos a Hugging Face
try:
# Subir el archivo .onnx
upload_file(
path_or_fileobj=onnx_file,
path_in_repo=f"{model_name}.onnx",
repo_id=repo_id,
token=token,
)
print(f"{model_name}.onnx subido a Hugging Face.")
# Subir el archivo .onnx.json
upload_file(
path_or_fileobj=onnx_json_file,
path_in_repo=f"{model_name}.onnx.json",
repo_id=repo_id,
token=token,
)
print(f"{model_name}.onnx.json subido a Hugging Face.")
except Exception as e:
print(f"Error al subir archivos a Hugging Face: {e}")
if __name__ == '__main__':
main()
|