from huggingface_hub import HfApi from fastapi import FastAPI, HTTPException from pydantic import BaseModel import requests import boto3 from dotenv import load_dotenv import os import uvicorn from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer import safetensors.torch from fastapi.responses import StreamingResponse from tqdm import tqdm # Cargar las variables de entorno desde el archivo .env load_dotenv() # Cargar las credenciales de AWS y el token de Hugging Face desde las variables de entorno AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID") AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY") AWS_REGION = os.getenv("AWS_REGION") S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME") # Nombre del bucket de S3 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN") # Token de Hugging Face # Cliente S3 de Amazon s3_client = boto3.client( 's3', aws_access_key_id=AWS_ACCESS_KEY_ID, aws_secret_access_key=AWS_SECRET_ACCESS_KEY, region_name=AWS_REGION ) app = FastAPI() # Pydantic Model para el cuerpo de la solicitud del endpoint /download_model/ class DownloadModelRequest(BaseModel): model_name: str pipeline_task: str input_text: str revision: str = "main" # Revisión por defecto class S3DirectStream: def __init__(self, bucket_name): self.s3_client = boto3.client( 's3', aws_access_key_id=AWS_ACCESS_KEY_ID, aws_secret_access_key=AWS_SECRET_ACCESS_KEY, region_name=AWS_REGION ) self.bucket_name = bucket_name def stream_from_s3(self, key): try: print(f"Descargando archivo {key} desde S3...") response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key) return response['Body'] # Devolver el cuerpo directamente para el StreamingResponse except self.s3_client.exceptions.NoSuchKey: raise HTTPException(status_code=404, detail=f"El archivo {key} no existe en el bucket S3.") def file_exists_in_s3(self, key): try: self.s3_client.head_object(Bucket=self.bucket_name, Key=key) return True except self.s3_client.exceptions.ClientError: return False def load_model_from_stream(self, model_prefix, revision): try: print(f"Cargando el modelo {model_prefix} desde S3...") if self.file_exists_in_s3(f"{model_prefix}/config.json") and \ (self.file_exists_in_s3(f"{model_prefix}/pytorch_model.bin") or self.file_exists_in_s3(f"{model_prefix}/model.safetensors")): print(f"Modelo {model_prefix} ya existe en S3. No es necesario descargarlo.") return self.load_model_from_existing_s3(model_prefix) print(f"Modelo {model_prefix} no encontrado. Procediendo a descargar...") self.download_and_upload_to_s3(model_prefix, revision) # Pasamos 'revision' aquí return self.load_model_from_stream(model_prefix, revision) except HTTPException as e: print(f"Error al cargar el modelo: {e}") return None def load_model_from_existing_s3(self, model_prefix): # Cargar el modelo y los archivos necesarios desde S3 print(f"Cargando los archivos {model_prefix} desde S3...") config_stream = self.stream_from_s3(f"{model_prefix}/config.json") config_data = config_stream.read().decode("utf-8") print(f"Cargando el modelo de lenguaje {model_prefix}...") # Verificar si el archivo es un safetensor o un archivo binario if self.file_exists_in_s3(f"{model_prefix}/model.safetensors"): # Usar safetensors si el archivo es de tipo safetensors model_stream = self.stream_from_s3(f"{model_prefix}/model.safetensors") model = AutoModelForCausalLM.from_config(config_data) model.load_state_dict(safetensors.torch.load_stream(model_stream)) # Cargar el modelo utilizando safetensors else: # Cargar el modelo utilizando pytorch si el archivo es .bin model_stream = self.stream_from_s3(f"{model_prefix}/pytorch_model.bin") model = AutoModelForCausalLM.from_config(config_data) model.load_state_dict(torch.load(model_stream, map_location="cpu")) return model def load_tokenizer_from_stream(self, model_prefix): try: if self.file_exists_in_s3(f"{model_prefix}/tokenizer.json"): print(f"Tokenizer para {model_prefix} ya existe en S3. No es necesario descargarlo.") return self.load_tokenizer_from_existing_s3(model_prefix) print(f"Tokenizer para {model_prefix} no encontrado. Procediendo a descargar...") self.download_and_upload_to_s3(model_prefix) # Pasamos 'revision' aquí también return self.load_tokenizer_from_stream(model_prefix) except HTTPException as e: print(f"Error al cargar el tokenizer: {e}") return None def load_tokenizer_from_existing_s3(self, model_prefix): print(f"Cargando el tokenizer para {model_prefix} desde S3...") tokenizer_stream = self.stream_from_s3(f"{model_prefix}/tokenizer.json") tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream) return tokenizer def download_and_upload_to_s3(self, model_prefix, revision): # URLs de los archivos de Hugging Face model_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/pytorch_model.bin" safetensors_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/model.safetensors" tokenizer_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/tokenizer.json" config_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/config.json" print(f"Descargando y subiendo archivos para el modelo {model_prefix} a S3...") self.download_and_upload_to_s3_url(model_url, f"{model_prefix}/pytorch_model.bin") self.download_and_upload_to_s3_url(safetensors_url, f"{model_prefix}/model.safetensors") self.download_and_upload_to_s3_url(tokenizer_url, f"{model_prefix}/tokenizer.json") self.download_and_upload_to_s3_url(config_url, f"{model_prefix}/config.json") def download_and_upload_to_s3_url(self, url: str, s3_key: str): print(f"Descargando archivo desde {url}...") response = requests.get(url) if response.status_code == 200: # Subir archivo a S3 print(f"Subiendo archivo a S3 con key {s3_key}...") self.s3_client.put_object(Bucket=self.bucket_name, Key=s3_key, Body=response.content) else: raise HTTPException(status_code=500, detail=f"Error al descargar el archivo desde {url}") @app.post("/predict/") async def predict(model_request: DownloadModelRequest): try: print(f"Recibiendo solicitud para predecir con el modelo {model_request.model_name}...") model_name = model_request.model_name revision = model_request.revision # Cargar el modelo y tokenizer desde S3 streamer = S3DirectStream(S3_BUCKET_NAME) model = streamer.load_model_from_stream(model_name, revision) tokenizer = streamer.load_tokenizer_from_stream(model_name) # Obtener el pipeline adecuado según la solicitud task = model_request.pipeline_task if task not in ["text-generation", "sentiment-analysis", "translation", "fill-mask", "question-answering", "text-to-speech", "text-to-image", "text-to-audio", "text-to-video"]: raise HTTPException(status_code=400, detail="Pipeline task no soportado") # Crear el pipeline dinámicamente basado en el tipo de tarea nlp_pipeline = pipeline(task, model=model, tokenizer=tokenizer, use_auth_token=HUGGINGFACE_TOKEN, revision=revision) # Ejecutar el pipeline con el input_text outputs = nlp_pipeline(model_request.input_text) # Almacenar el resultado en S3 dependiendo del tipo de tarea if task == "text-to-image": s3_key = f"{model_request.model_name}/generated_image.png" return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="image/png") elif task == "text-to-speech": s3_key = f"{model_request.model_name}/generated_audio.wav" return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="audio/wav") elif task == "text-to-video": s3_key = f"{model_request.model_name}/generated_video.mp4" return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="video/mp4") # Devolver resultados de texto u otros tipos de tarea return {"result": outputs} except Exception as e: raise HTTPException(status_code=500, detail=f"Error al procesar la solicitud: {str(e)}") if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=8000)