Update app.py

app.py

@@ -4,23 +4,23 @@ import torch
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import safetensors.torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from fastapi.responses import StreamingResponse
 import io
-from tqdm import tqdm
 import requests
 import uvicorn
+from dotenv import load_dotenv
 import re
-import sys
+from tqdm import tqdm
+
 # Cargar las variables de entorno desde el archivo .env
-from dotenv import load_dotenv
 load_dotenv()
 
-# Cargar las credenciales de AWS desde las variables de entorno
+# 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 de S3
+S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")  # Nombre del bucket S3
 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")  # Token de Hugging Face
 
 # Cliente S3 de Amazon
@@ -67,31 +67,32 @@ class S3DirectStream:
         except self.s3_client.exceptions.ClientError:
             return False
 
-    def load_model_from_stream(self, model_prefix):
+    def load_model_from_stream(self, model_name):
         try:
-            print(f"[INFO] Cargando el modelo {model_prefix} desde S3...")
+            print(f"[INFO] Cargando el modelo {model_name} desde S3...")
+            model_prefix = model_name.lower()
             model_files = self.get_model_file_parts(model_prefix)
             if not model_files:
-                model_files = [f"{model_prefix}/model"]  # Uso de modelo base
-
+                print(f"[INFO] Modelo no encontrado en S3, descargando desde Hugging Face a S3...")
+                self.download_and_upload_from_huggingface(model_name)
+                model_files = self.get_model_file_parts(model_prefix)
+            
             # Leer y cargar todos los archivos del modelo
             model_streams = []
             for model_file in tqdm(model_files, desc="Cargando archivos del modelo", unit="archivo"):
                 model_streams.append(self.stream_from_s3(model_file))
 
-            # Verificar si el archivo es un safetensor o un archivo binario
             config_stream = self.stream_from_s3(f"{model_prefix}/config.json")
             config_data = config_stream.read().decode("utf-8")
             
-            # Cargar el modelo dependiendo de si es safetensor o binario
             if model_files[0].endswith("model.safetensors"):
                 print("[INFO] Cargando el modelo como safetensor...")
                 model = AutoModelForCausalLM.from_config(config_data)
-                model.load_state_dict(safetensors.torch.load_stream(model_streams[0]))  # Cargar el modelo utilizando safetensors
+                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"))  # Cargar el modelo utilizando pytorch
+                model.load_state_dict(torch.load(model_streams[0], map_location="cpu"))
 
             print("[INFO] Modelo cargado con éxito.")
             return model
@@ -100,28 +101,46 @@ class S3DirectStream:
             print(f"[ERROR] Error al cargar el modelo desde S3: {e}")
             raise HTTPException(status_code=500, detail="Error al cargar el modelo desde S3.")
 
-    def load_tokenizer_from_stream(self, model_prefix):
+    def load_tokenizer_from_stream(self, model_name):
         try:
-            print(f"[INFO] Cargando el tokenizer {model_prefix} desde S3...")
-            tokenizer_stream = self.stream_from_s3(f"{model_prefix}/tokenizer.json")
+            print(f"[INFO] Cargando el tokenizer {model_name} desde S3...")
+            tokenizer_stream = self.stream_from_s3(f"{model_name}/tokenizer.json")
             tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
             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_prefix):
-        print(f"[INFO] Listando archivos del modelo en S3 con prefijo {model_prefix}...")
-        files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_prefix)
+    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(files.get('Contents', []), desc="Verificando archivos", unit="archivo"):
             key = obj['Key']
-            if re.match(rf"{model_prefix}/model(-\d+-of-\d+)?", key) or key.endswith("model.safetensors"):
+            if re.match(rf"{model_name}/model(-\d+-of-\d+)?", key) or key.endswith("model.safetensors"):
                 model_files.append(key)
         if not model_files:
-            print(f"[WARNING] No se encontraron archivos coincidentes con el patrón para el modelo {model_prefix}.")
+            print(f"[WARNING] No se encontraron archivos para el modelo {model_name}.")
         return model_files
 
+    def download_and_upload_from_huggingface(self, model_name):
+        try:
+            print(f"[INFO] Descargando {model_name} desde Hugging Face...")
+            model_url = f"https://huggingface.co/{model_name}/resolve/main/"
+            files_to_download = [
+                f"{model_name}/pytorch_model.bin",
+                f"{model_name}/config.json",
+                f"{model_name}/tokenizer.json",
+            ]
+            
+            for file in files_to_download:
+                file_url = model_url + file
+                s3_key = file
+                self.download_and_upload_to_s3_url(file_url, 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.")
+
     def download_and_upload_to_s3_url(self, url: str, s3_key: str):
         try:
             print(f"[INFO] Descargando archivo desde {url}...")
@@ -129,8 +148,6 @@ class S3DirectStream:
             if response.status_code == 200:
                 print(f"[INFO] Subiendo archivo a S3 con key {s3_key}...")
                 self.s3_client.put_object(Bucket=self.bucket_name, Key=s3_key, Body=response.content)
-                # Eliminar el archivo local después de la carga exitosa
-                self.delete_local_file(s3_key)
             else:
                 print(f"[ERROR] Error al descargar el archivo desde {url}, código de estado {response.status_code}.")
                 raise HTTPException(status_code=500, detail=f"Error al descargar el archivo desde {url}")
@@ -138,18 +155,6 @@ class S3DirectStream:
             print(f"[ERROR] Error al procesar la URL {url}: {str(e)}")
             raise HTTPException(status_code=500, detail=f"Error al procesar la URL {url}")
 
-    def delete_local_file(self, file_path: str):
-        try:
-            print(f"[INFO] Eliminando archivo local {file_path}...")
-            if os.path.exists(file_path):
-                os.remove(file_path)
-                print(f"[INFO] Archivo local {file_path} eliminado correctamente.")
-            else:
-                print(f"[WARNING] El archivo local {file_path} no existe.")
-        except Exception as e:
-            print(f"[ERROR] Error al eliminar el archivo local {file_path}: {str(e)}")
-
-
 @app.post("/predict/")
 async def predict(model_request: DownloadModelRequest):
     try:
@@ -188,4 +193,4 @@ async def predict(model_request: DownloadModelRequest):
 # Iniciar servidor de predicciones
 if __name__ == "__main__":
     print("Iniciando servidor de predicciones en localhost:8000")
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=8000)