Update app.py

app.py

@@ -10,16 +10,18 @@ import torch
 import safetensors.torch
 from fastapi.responses import StreamingResponse
 from tqdm import tqdm
-import re
 
+# 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")
-HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+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,
@@ -29,6 +31,7 @@ s3_client = boto3.client(
 
 app = FastAPI()
 
+# Pydantic Model para el cuerpo de la solicitud del endpoint /download_model/
 class DownloadModelRequest(BaseModel):
     model_name: str
     pipeline_task: str
@@ -46,135 +49,145 @@ class S3DirectStream:
 
     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']
+            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.")
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al descargar de S3: {e}")
 
     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 as e:
-            if e.response['Error']['Code'] == '404':
-                return False
-            raise HTTPException(status_code=500, detail=f"Error al verificar archivo en S3: {e}")
+        except self.s3_client.exceptions.ClientError:
+            return False
 
     def load_model_from_stream(self, model_prefix):
         try:
-            model_files = self.list_model_files(model_prefix)
-            if not model_files:
-                self.download_and_upload_to_s3(model_prefix)
-                return self.load_model_from_stream(model_prefix)
-
-            config_stream = self.stream_from_s3(f"{model_prefix}/config.json")
-            config_data = config_stream.read().decode("utf-8")
-
-            model_path = f"{model_prefix}/model.safetensors"
-            if self.file_exists_in_s3(model_path):
-                model_stream = self.stream_from_s3(model_path)
-                model = AutoModelForCausalLM.from_config(config_data)
-                model.load_state_dict(safetensors.torch.load_stream(model_stream))
-            elif model_files:
-                model = AutoModelForCausalLM.from_config(config_data)
-                state_dict = {}
-                for file_name in model_files:
-                    file_stream = self.stream_from_s3(f"{model_prefix}/{file_name}")
-                    tmp = torch.load(file_stream, map_location="cpu")
-                    state_dict.update(tmp)
-                model.load_state_dict(state_dict)
-            else:
-                raise HTTPException(status_code=500, detail="Modelo no encontrado")
-
-            return model
+            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)
+            return self.load_model_from_stream(model_prefix)
         except HTTPException as e:
-            raise
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al cargar el modelo: {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"))
 
-    def list_model_files(self, model_prefix):
-        try:
-          response = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=f"{model_prefix}/pytorch_model-")
-          model_files = []
-          if 'Contents' in response:
-            for obj in response['Contents']:
-              if re.match(r"pytorch_model-\d+-of-\d+", obj['Key'].split('/')[-1]):
-                model_files.append(obj['Key'].split('/')[-1])
-          return model_files
-        except Exception as e:
-          return None
+        return model
 
     def load_tokenizer_from_stream(self, model_prefix):
         try:
-            tokenizer_path = f"{model_prefix}/tokenizer.json"
-            if self.file_exists_in_s3(tokenizer_path):
-                tokenizer_stream = self.stream_from_s3(tokenizer_path)
-                tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
-                return tokenizer
-            else:
-                self.download_and_upload_to_s3(model_prefix)
-                return self.load_tokenizer_from_stream(model_prefix)
+            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)
+            return self.load_tokenizer_from_stream(model_prefix)
         except HTTPException as e:
-            raise
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al cargar el tokenizer: {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):
-        urls = {
-            "pytorch_model.bin": f"https://huggingface.co/{model_prefix}/resolve/main/pytorch_model.bin",
-            "model.safetensors": f"https://huggingface.co/{model_prefix}/resolve/main/model.safetensors",
-            "tokenizer.json": f"https://huggingface.co/{model_prefix}/resolve/main/tokenizer.json",
-            "config.json": f"https://huggingface.co/{model_prefix}/resolve/main/config.json"
-        }
-
-        for filename, url in urls.items():
-            try:
-                response = requests.get(url, stream=True)
-                response.raise_for_status()
-                self.s3_client.upload_fileobj(response.raw, self.bucket_name, f"{model_prefix}/{filename}")
-            except requests.exceptions.RequestException as e:
-                raise HTTPException(status_code=500, detail=f"Error al descargar {filename}: {e}")
-            except Exception as e:
-                raise HTTPException(status_code=500, detail=f"Error al subir {filename} a S3: {e}")
+        # URLs de los archivos de Hugging Face
+        model_url = f"https://huggingface.co/{model_prefix}/resolve/main/pytorch_model.bin"
+        safetensors_url = f"https://huggingface.co/{model_prefix}/resolve/main/model.safetensors"
+        tokenizer_url = f"https://huggingface.co/{model_prefix}/resolve/main/tokenizer.json"
+        config_url = f"https://huggingface.co/{model_prefix}/resolve/main/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}...")
+        # Cargar el modelo y tokenizer desde S3
         streamer = S3DirectStream(S3_BUCKET_NAME)
         model = streamer.load_model_from_stream(model_request.model_name)
         tokenizer = streamer.load_tokenizer_from_stream(model_request.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)
+
+        # Ejecutar el pipeline con el input_text
         input_text = model_request.input_text
         outputs = nlp_pipeline(input_text)
 
+        # Procesar los diferentes tipos de respuestas según el pipeline
         if task in ["text-generation", "translation", "fill-mask", "sentiment-analysis", "question-answering"]:
             return {"response": outputs}
+
         elif task == "text-to-image":
-            s3_key = f"{model_request.model_name}/generated_image.png"
+            # Asumir que outputs es la imagen generada
+            s3_key = f"{model_request.model_name}/generated_image.png"  # Definir el key del archivo de imagen
             return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="image/png")
+
         elif task == "text-to-audio":
-            s3_key = f"{model_request.model_name}/generated_audio.wav"
+            # Asumir que outputs es el audio generado
+            s3_key = f"{model_request.model_name}/generated_audio.wav"  # Definir el key del archivo de audio
             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"
+            # Asumir que outputs es el video generado
+            s3_key = f"{model_request.model_name}/generated_video.mp4"  # Definir el key del archivo de video
             return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="video/mp4")
+
         else:
             raise HTTPException(status_code=400, detail="Tipo de tarea desconocido")
 
-    except HTTPException as e:
-        raise
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error inesperado: {str(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=7860)
+    uvicorn.run(app, host="0.0.0.0", port=8000)