Update app.py

app.py

@@ -1,128 +1,123 @@
 import os
-import json
-import requests
+import logging
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from google.cloud import storage
-from google.auth import exceptions
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from io import BytesIO
-from dotenv import load_dotenv
-import uvicorn
-import tempfile
+from transformers import pipeline
+import json
+from google.auth.exceptions import DefaultCredentialsError
 
-load_dotenv()
+# Configuración de GCS
 
+# Cargar las variables de entorno
 API_KEY = os.getenv("API_KEY")
 GCS_BUCKET_NAME = os.getenv("GCS_BUCKET_NAME")
 GOOGLE_APPLICATION_CREDENTIALS_JSON = os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON")
 HF_API_TOKEN = os.getenv("HF_API_TOKEN")
 
+# Configuración de logs
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
 try:
-    credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
-    storage_client = storage.Client.from_service_account_info(credentials_info)
-    bucket = storage_client.bucket(GCS_BUCKET_NAME)
-except (exceptions.DefaultCredentialsError, json.JSONDecodeError, KeyError, ValueError) as e:
-    raise RuntimeError(f"Error al cargar credenciales o bucket: {e}")
+    # Intentar cargar las credenciales de servicio de GCS desde la variable de entorno
+    credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)  # Cargar el JSON de credenciales
+    storage_client = storage.Client.from_service_account_info(credentials_info)  # Crear cliente de GCS
+    bucket = storage_client.bucket(GCS_BUCKET_NAME)  # Acceder al bucket
+
+    # Verificación exitosa
+    logger.info(f"Conexión con Google Cloud Storage exitosa. Bucket: {GCS_BUCKET_NAME}")
+
+except (DefaultCredentialsError, json.JSONDecodeError, KeyError, ValueError) as e:
+    # Manejo de errores en caso de que las credenciales sean incorrectas o faltantes
+    logger.error(f"Error al cargar las credenciales o bucket: {e}")
+    raise RuntimeError(f"Error al cargar las credenciales o bucket: {e}")
 
+# Configurar la aplicación FastAPI
 app = FastAPI()
 
-class DownloadModelRequest(BaseModel):
+# Configuración de logs
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+class PredictionRequest(BaseModel):
     model_name: str
     pipeline_task: str
     input_text: str
 
-class GCSHandler:
-    def __init__(self, bucket_name):
-        self.bucket = storage_client.bucket(bucket_name)
-
-    def file_exists(self, blob_name):
-        return self.bucket.blob(blob_name).exists()
+# Función para obtener la URL del modelo desde GCS
+def get_gcs_model_url(bucket_name: str, model_name: str):
+    """
+    Obtiene la URL del modelo desde GCS.
+    """
+    try:
+        model_dir = f"models/{model_name}/"
+        
+        # Verificar si la carpeta del modelo existe en GCS
+        bucket = storage_client.get_bucket(bucket_name)
+        blobs = bucket.list_blobs(prefix=model_dir)
 
-    def upload_file(self, blob_name, file_stream):
-        blob = self.bucket.blob(blob_name)
-        blob.upload_from_file(file_stream)
+        # Verificar si existen archivos en el directorio del modelo
+        file_list = [blob.name for blob in blobs]
+        if not file_list:
+            raise HTTPException(status_code=404, detail="No se encontraron los archivos del modelo en GCS.")
+        
+        # Construir la URL GCS del modelo (en este caso solo la ruta del directorio)
+        gcs_url = f"gs://{bucket_name}/{model_dir}"
+        
+        return gcs_url
 
-    def download_file(self, blob_name):
-        blob = self.bucket.blob(blob_name)
-        if not blob.exists():
-            raise HTTPException(status_code=404, detail=f"File '{blob_name}' not found.")
-        return BytesIO(blob.download_as_bytes())
+    except Exception as e:
+        logger.error(f"Error al obtener la URL del modelo desde GCS: {str(e)}")
+        raise HTTPException(status_code=500, detail="Error al obtener la URL del modelo desde GCS.")
 
-def download_model_from_huggingface(model_name):
-    url = f"https://huggingface.co/{model_name}/tree/main"
-    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
-    
+# Función para cargar el pipeline directamente desde GCS como URL
+def load_pipeline_from_gcs(model_name: str, pipeline_task: str):
+    """
+    Carga el pipeline directamente desde la URL del modelo en GCS sin usar RAM ni almacenamiento temporal.
+    """
     try:
-        response = requests.get(url, headers=headers)
-        if response.status_code == 200:
-            model_files = [
-                "pytorch_model.bin",
-                "config.json",
-                "tokenizer.json",
-                "model.safetensors",
-            ]
-            for file_name in model_files:
-                file_url = f"https://huggingface.co/{model_name}/resolve/main/{file_name}"
-                file_content = requests.get(file_url).content
-                blob_name = f"{model_name}/{file_name}"
-                bucket.blob(blob_name).upload_from_file(BytesIO(file_content))
-        else:
-            raise HTTPException(status_code=404, detail="Error al acceder al árbol de archivos de Hugging Face.")
+        # Obtener la URL del modelo desde GCS
+        model_url = get_gcs_model_url(GCS_BUCKET_NAME, model_name)
+        
+        # Cargar el pipeline directamente desde la URL del modelo
+        nlp_pipeline = pipeline(
+            task=pipeline_task,
+            model=model_url,  # Usamos la URL de GCS como modelo
+        )
+
+        return nlp_pipeline
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error descargando archivos de Hugging Face: {e}")
+        logger.error(f"Error al cargar el pipeline desde GCS: {str(e)}")
+        raise HTTPException(status_code=500, detail="Error al cargar el pipeline desde GCS.")
 
-@app.post("/predict/")
-async def predict(request: DownloadModelRequest):
+# Endpoint para realizar la predicción
+@app.post("/predict")
+def predict(request: PredictionRequest):
+    """
+    Endpoint para recibir solicitudes POST con datos JSON y realizar la predicción.
+    """
     try:
-        gcs_handler = GCSHandler(GCS_BUCKET_NAME)
-        model_prefix = request.model_name
-        model_files = [
-            "pytorch_model.bin",
-            "config.json",
-            "tokenizer.json",
-            "model.safetensors",
-        ]
-        
-        model_files_exist = all(gcs_handler.file_exists(f"{model_prefix}/{file}") for file in model_files)
+        # Extraer los parámetros de la solicitud JSON
+        model_name = request.model_name
+        pipeline_task = request.pipeline_task
+        input_text = request.input_text
         
-        if not model_files_exist:
-            download_model_from_huggingface(model_prefix)
+        # Cargar el pipeline directamente desde GCS sin usar RAM ni almacenamiento temporal
+        nlp_pipeline = load_pipeline_from_gcs(model_name, pipeline_task)
         
-        model_files_streams = {file: gcs_handler.download_file(f"{model_prefix}/{file}") for file in model_files if gcs_handler.file_exists(f"{model_prefix}/{file}")}
-        
-        config_stream = model_files_streams.get("config.json")
-        tokenizer_stream = model_files_streams.get("tokenizer.json")
-        model_stream = model_files_streams.get("pytorch_model.bin")
-        
-        if not config_stream or not tokenizer_stream or not model_stream:
-            raise HTTPException(status_code=500, detail="Required model files missing.")
-        
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            config_path = os.path.join(tmp_dir, "config.json")
-            tokenizer_path = os.path.join(tmp_dir, "tokenizer.json")
-            model_path = os.path.join(tmp_dir, "pytorch_model.bin")
-
-            with open(config_path, 'wb') as f:
-                f.write(config_stream.read())
-            with open(tokenizer_path, 'wb') as f:
-                f.write(tokenizer_stream.read())
-            with open(model_path, 'wb') as f:
-                f.write(model_stream.read())
-            
-            model = AutoModelForCausalLM.from_pretrained(tmp_dir, from_tf=True)
-            tokenizer = AutoTokenizer.from_pretrained(tmp_dir)
-            
-            pipeline_ = pipeline(request.pipeline_task, model=model, tokenizer=tokenizer)
-            
-            result = pipeline_(request.input_text)
+        # Realizar la predicción
+        result = nlp_pipeline(input_text)
         
         return {"response": result}
-    
+
     except HTTPException as e:
+        logger.error(f"Error en la predicción: {e.detail}")
         raise e
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error: {e}")
+        logger.error(f"Error en la predicción: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
+    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)