Create app.py

app.py

@@ -0,0 +1,191 @@
+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)