Update app.py

app.py

@@ -1,29 +1,23 @@
 import os
 import json
+import logging
 import boto3
-import uvicorn
 from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
+from fastapi.responses import JSONResponse
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from huggingface_hub import hf_hub_download
-from io import BytesIO
-import torch
-import safetensors
-from dotenv import load_dotenv
-import tqdm
-import re
+from tqdm import tqdm
+import io
 
-# Cargar las variables de entorno desde el archivo .env
-load_dotenv()
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# Configuración de AWS y Hugging Face
 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 S3
-HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")  # Token de Hugging Face
+S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
+HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
 
-# Cliente S3 de Amazon
 s3_client = boto3.client(
     's3',
     aws_access_key_id=AWS_ACCESS_KEY_ID,
@@ -33,11 +27,16 @@ s3_client = boto3.client(
 
 app = FastAPI()
 
-# Pydantic Model para el cuerpo de la solicitud del endpoint /predict/
-class DownloadModelRequest(BaseModel):
-    model_name: str
-    pipeline_task: str
-    input_text: str
+PIPELINE_MAP = {
+    "text-generation": "text-generation",
+    "sentiment-analysis": "sentiment-analysis",
+    "translation": "translation",
+    "fill-mask": "fill-mask",
+    "question-answering": "question-answering",
+    "text-to-speech": "text-to-speech",
+    "text-to-video": "text-to-video",
+    "text-to-image": "text-to-image"
+}
 
 class S3DirectStream:
     def __init__(self, bucket_name):
@@ -51,155 +50,123 @@ class S3DirectStream:
 
     def stream_from_s3(self, key):
         try:
-            print(f"[INFO] Descargando archivo {key} desde S3...")
+            logger.info(f"Descargando {key} desde S3...")
             response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
-            return BytesIO(response['Body'].read())  # Devolver el cuerpo como BytesIO
+            return response['Body']
         except self.s3_client.exceptions.NoSuchKey:
+            logger.error(f"El archivo {key} no existe en el bucket S3.")
             raise HTTPException(status_code=404, detail=f"El archivo {key} no existe en el bucket S3.")
         except Exception as e:
-            print(f"[ERROR] Error al descargar {key}: {str(e)}")
-            raise HTTPException(status_code=500, detail=f"Error al descargar archivo {key} desde S3.")
+            logger.error(f"Error al descargar {key} desde S3: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"Error al descargar {key} desde S3: {str(e)}")
 
-    def file_exists_in_s3(self, key):
+    def get_model_file_parts(self, model_name):
         try:
-            self.s3_client.head_object(Bucket=self.bucket_name, Key=key)
-            return True
-        except self.s3_client.exceptions.ClientError:
-            return False
+            model_prefix = model_name.lower()
+            logger.info(f"Obteniendo archivos para el modelo {model_name} desde S3...")
+            files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_prefix)
+            model_files = [obj['Key'] for obj in files.get('Contents', []) if model_prefix in obj['Key']]
+            return model_files
+        except Exception as e:
+            logger.error(f"Error al obtener archivos del modelo {model_name} desde S3: {e}")
+            raise HTTPException(status_code=500, detail=f"Error al obtener archivos del modelo {model_name} desde S3: {e}")
 
     def load_model_from_s3(self, model_name):
         try:
-            print(f"[INFO] Cargando el modelo {model_name} desde S3...")
-
-            # Verificar si el modelo existe en S3
             model_prefix = model_name.lower()
             model_files = self.get_model_file_parts(model_prefix)
 
             if not model_files:
-                print(f"[INFO] El modelo {model_name} no está en S3. Procediendo a descargar desde Hugging Face...")
+                logger.info(f"El modelo {model_name} no está en S3, descargando desde Hugging Face...")
                 self.download_and_upload_from_huggingface(model_name)
                 model_files = self.get_model_file_parts(model_prefix)
 
             if not model_files:
+                logger.error(f"Archivos del modelo {model_name} no encontrados en S3.")
                 raise HTTPException(status_code=404, detail=f"Archivos del modelo {model_name} no encontrados en S3.")
 
-            # Cargar todos los archivos del modelo desde S3
-            model_streams = []
-            for model_file in tqdm.tqdm(model_files, desc="Cargando archivos del modelo", unit="archivo"):
-                model_streams.append(self.stream_from_s3(model_file))
-
+            logger.info(f"Cargando archivos del modelo {model_name}...")
             config_stream = self.stream_from_s3(f"{model_prefix}/config.json")
-            config_data = json.loads(config_stream.read().decode("utf-8"))
-
-            # Cargar el modelo dependiendo del tipo de archivo (torch o safetensors)
-            if any(file.endswith("model.safetensors") for file in model_files):
-                print("[INFO] Cargando el modelo como safetensor...")
-                model = AutoModelForCausalLM.from_config(config_data)
-                model.load_state_dict(safetensors.torch.load_stream(model_streams[0]))
-            else:
-                print("[INFO] Cargando el modelo como archivo binario de PyTorch...")
-                model = AutoModelForCausalLM.from_config(config_data)
-                model.load_state_dict(torch.load(model_streams[0], map_location="cpu"))
-
-            print("[INFO] Modelo cargado con éxito.")
+            config_data = config_stream.read()
+
+            if not config_data:
+                logger.error(f"El archivo de configuración {model_prefix}/config.json está vacío.")
+                raise HTTPException(status_code=500, detail=f"El archivo de configuración {model_prefix}/config.json está vacío.")
+            
+            config_text = config_data.decode("utf-8")
+            config_json = json.loads(config_text)
+
+            model = AutoModelForCausalLM.from_pretrained(f"s3://{self.bucket_name}/{model_prefix}", config=config_json, from_tf=False)
             return model
 
         except Exception as e:
-            print(f"[ERROR] Error al cargar el modelo desde S3: {e}")
-            raise HTTPException(status_code=500, detail="Error al cargar el modelo desde S3.")
+            logger.error(f"Error al cargar el modelo desde S3: {e}")
+            raise HTTPException(status_code=500, detail=f"Error al cargar el modelo desde S3: {e}")
 
     def load_tokenizer_from_s3(self, model_name):
         try:
-            print(f"[INFO] Cargando el tokenizer {model_name} desde S3...")
+            logger.info(f"Cargando el tokenizer del modelo {model_name} desde S3...")
             tokenizer_stream = self.stream_from_s3(f"{model_name}/tokenizer.json")
-            tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
+            tokenizer_data = tokenizer_stream.read().decode("utf-8")
+            tokenizer = AutoTokenizer.from_pretrained(f"s3://{self.bucket_name}/{model_name}")
             return tokenizer
         except Exception as e:
-            print(f"[ERROR] Error al cargar el tokenizer desde S3: {e}")
-            raise HTTPException(status_code=500, detail="Error al cargar el tokenizer desde S3.")
-
-    def get_model_file_parts(self, model_name):
-        print(f"[INFO] Listando archivos del modelo en S3 con prefijo {model_name}...")
-        files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_name)
-        model_files = []
-        for obj in tqdm.tqdm(files.get('Contents', []), desc="Verificando archivos", unit="archivo"):
-            key = obj['Key']
-            # Verificar si es un archivo relevante del modelo
-            if re.match(rf"{model_name}/.*", key):
-                model_files.append(key)
-        if not model_files:
-            print(f"[WARNING] No se encontraron archivos para el modelo {model_name}.")
-        return model_files
+            logger.error(f"Error al cargar el tokenizer desde S3: {e}")
+            raise HTTPException(status_code=500, detail=f"Error al cargar el tokenizer desde S3: {e}")
 
     def download_and_upload_from_huggingface(self, model_name):
         try:
-            print(f"[INFO] Descargando {model_name} desde Hugging Face...")
-            # Descargar todos los archivos del modelo
-            files_to_download = hf_hub_download(repo_id=model_name, use_auth_token=HUGGINGFACE_TOKEN)
+            logger.info(f"Descargando modelo {model_name} desde Hugging Face...")
+            files_to_download = hf_hub_download(repo_id=model_name, use_auth_token=HUGGINGFACE_TOKEN, local_dir=model_name)
 
-            # Subir a S3 los archivos descargados
-            for file in files_to_download:
+            for file in tqdm(files_to_download, desc="Subiendo archivos a S3"):
                 file_name = os.path.basename(file)
                 s3_key = f"{model_name}/{file_name}"
                 if not self.file_exists_in_s3(s3_key):
                     self.upload_file_to_s3(file, s3_key)
 
         except Exception as e:
-            print(f"[ERROR] Error al descargar y subir modelo desde Hugging Face: {e}")
-            raise HTTPException(status_code=500, detail="Error al descargar y subir modelo desde Hugging Face.")
+            logger.error(f"Error al descargar y subir modelo desde Hugging Face: {e}")
+            raise HTTPException(status_code=500, detail=f"Error al descargar y subir modelo desde Hugging Face: {e}")
 
     def upload_file_to_s3(self, file_path, s3_key):
         try:
-            print(f"[INFO] Subiendo archivo {file_path} a S3 con key {s3_key}...")
             with open(file_path, 'rb') as data:
                 self.s3_client.put_object(Bucket=self.bucket_name, Key=s3_key, Body=data)
-            # Eliminar archivo local después de ser subido
             os.remove(file_path)
+            logger.info(f"Archivo {file_path} subido correctamente a S3 y eliminado localmente.")
         except Exception as e:
-            print(f"[ERROR] Error al subir archivo a S3: {e}")
-            raise HTTPException(status_code=500, detail="Error al subir archivo a S3.")
-
+            logger.error(f"Error al subir archivo a S3: {e}")
+            raise HTTPException(status_code=500, detail=f"Error al subir archivo a S3: {e}")
 
 @app.post("/predict/")
-async def predict(model_request: DownloadModelRequest):
+async def predict(model_request: dict):
     try:
-        print(f"[INFO] Recibiendo solicitud para predecir con el modelo {model_request.model_name}...")
+        model_name = model_request.get("model_name")
+        task = model_request.get("pipeline_task")
+        input_text = model_request.get("input_text")
+
         streamer = S3DirectStream(S3_BUCKET_NAME)
-        model = streamer.load_model_from_s3(model_request.model_name)
-        tokenizer = streamer.load_tokenizer_from_s3(model_request.model_name)
+        model = streamer.load_model_from_s3(model_name)
+        tokenizer = streamer.load_tokenizer_from_s3(model_name)
 
-        task = model_request.pipeline_task
-        if task not in ["text-generation", "sentiment-analysis", "translation", "fill-mask", "question-answering", 
-                        "text-to-speech", "text-to-video", "text-to-image"]:
+        if task not in PIPELINE_MAP:
+            logger.error("Pipeline task no soportado")
             raise HTTPException(status_code=400, detail="Pipeline task no soportado")
 
-        # Configurar el pipeline de transformers según la tarea
-        nlp_pipeline = None
-        if task == "text-generation":
-            nlp_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
-        elif task == "sentiment-analysis":
-            nlp_pipeline = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer)
-        elif task == "translation":
-            nlp_pipeline = pipeline("translation", model=model, tokenizer=tokenizer)
-        elif task == "fill-mask":
-            nlp_pipeline = pipeline("fill-mask", model=model, tokenizer=tokenizer)
-        elif task == "question-answering":
-            nlp_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)
-        elif task == "text-to-speech":
-            nlp_pipeline = pipeline("text-to-speech", model=model, tokenizer=tokenizer)
-        elif task == "text-to-video":
-            nlp_pipeline = pipeline("text-to-video", model=model, tokenizer=tokenizer)
-        elif task == "text-to-image":
-            nlp_pipeline = pipeline("text-to-image", model=model, tokenizer=tokenizer)
-
-        result = nlp_pipeline(model_request.input_text)
-        return {"result": result}
+        nlp_pipeline = pipeline(PIPELINE_MAP[task], model=model, tokenizer=tokenizer)
 
-    except Exception as e:
-        print(f"[ERROR] Error en el proceso de predicción: {str(e)}")
-        raise HTTPException(status_code=500, detail="Error en el proceso de predicción")
+        result = nlp_pipeline(input_text)
+
+        if isinstance(result, dict) and 'file' in result:
+            return JSONResponse(content={"file": result['file']})
+        else:
+            return JSONResponse(content={"result": result})
 
+    except Exception as e:
+        logger.error(f"Error al realizar la predicción: {e}")
+        raise HTTPException(status_code=500, detail=f"Error al realizar la predicción: {e}")
 
-# Ejecutar la app con Uvicorn
 if __name__ == "__main__":
+    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)