Update app.py

app.py

@@ -1,108 +1,114 @@
 import os
 import logging
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from google.cloud import storage
+from transformers import pipeline
 import json
-from fastapi import FastAPI, HTTPException, Query
-from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
-import requests
-import uvicorn
-import threading
-from dotenv import load_dotenv
+from google.auth.exceptions import DefaultCredentialsError
 
-load_dotenv()
-
-logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
+# 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:
-    from google.cloud import storage
-    from google.auth import exceptions
-    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)
+    # 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 ModuleNotFoundError as e:
-    logger.error(f"Falta el módulo requerido: {str(e)}")
-    raise RuntimeError("El entorno no tiene instalado el paquete necesario 'google-cloud'.")
-except (exceptions.DefaultCredentialsError, json.JSONDecodeError, KeyError, ValueError) as e:
+
+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()
 
-HF_API_URL = "https://huggingface.co/api/models"
-
-def list_huggingface_models():
-    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
-    response = requests.get(HF_API_URL, headers=headers)
-    if response.status_code != 200:
-        logger.error(f"Error al obtener la lista de modelos de Hugging Face: HTTP {response.status_code}")
-        raise HTTPException(status_code=response.status_code, detail=f"Error al obtener la lista de modelos de Hugging Face")
-    return [model["modelId"] for model in response.json()]
-
-def download_model_to_bucket(model_name: str):
-    model_folder = f"{model_name}/"
-    if bucket.blob(f"{model_folder}config.json").exists():
-        logger.info(f"El modelo '{model_name}' ya está en el bucket.")
-        return
-
-    hf_base_url = f"https://huggingface.co/{model_name}/resolve/main/"
-    files = ["config.json", "pytorch_model.bin", "tokenizer.json"]
-
-    for file_name in files:
-        file_url = hf_base_url + file_name
-        blob = bucket.blob(f"{model_folder}{file_name}")
-        logger.info(f"Descargando {file_url} al bucket...")
-
-        response = requests.get(file_url, stream=True)
-        if response.status_code == 200:
-            blob.upload_from_string(response.content)
-            logger.info(f"Archivo {file_name} subido al bucket.")
-        else:
-            logger.error(f"No se pudo descargar el archivo {file_name}. HTTP {response.status_code}")
-
-def load_model_from_bucket(model_name: str):
-    model_folder = f"{model_name}/"
-    model_blob = bucket.blob(f"{model_folder}pytorch_model.bin")
-    config_blob = bucket.blob(f"{model_folder}config.json")
-    tokenizer_blob = bucket.blob(f"{model_folder}tokenizer.json")
-
-    if not (model_blob.exists() and config_blob.exists() and tokenizer_blob.exists()):
-        logger.error(f"Modelo no encontrado en el bucket: {model_name}")
-        raise HTTPException(status_code=404, detail="Modelo no encontrado en el bucket. Asegúrese de haberlo descargado correctamente.")
-
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_blob.download_as_bytes())
-    model = AutoModelForSequenceClassification.from_pretrained(model_blob.download_as_bytes(), config=config_blob.download_as_bytes())
-    return tokenizer, model
-
-def background_download_all_models():
-    logger.info("Iniciando descarga de modelos en segundo plano...")
+# 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
+
+# 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_names = list_huggingface_models()
-        for model_name in model_names:
-            try:
-                download_model_to_bucket(model_name)
-            except Exception as e:
-                logger.error(f"Error al descargar el modelo {model_name}: {str(e)}")
-    except HTTPException as e:
-        logger.error(f"Error al obtener la lista de modelos de Hugging Face: {str(e)}")
+        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)
+
+        # 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
 
-threading.Thread(target=background_download_all_models, daemon=True).start()
+    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.")
+
+# 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:
+        # 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:
+        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.")
 
+# Endpoint para realizar la predicción
 @app.post("/predict")
-def predict(
-    model_name: str = Query(..., description="Nombre del modelo"),
-    pipeline_task: str = Query(..., description="Tarea del pipeline (e.g., text-generation)"),
-    input_text: str = Query(..., description="Texto de entrada para el modelo")
-):
+def predict(request: PredictionRequest):
+    """
+    Endpoint para recibir solicitudes POST con datos JSON y realizar la predicción.
+    """
     try:
-        download_model_to_bucket(model_name)
-        tokenizer, model = load_model_from_bucket(model_name)
-        nlp_pipeline = pipeline(task=pipeline_task, model=model, tokenizer=tokenizer)
+        # Extraer los parámetros de la solicitud JSON
+        model_name = request.model_name
+        pipeline_task = request.pipeline_task
+        input_text = request.input_text
+        
+        # Cargar el pipeline directamente desde GCS sin usar RAM ni almacenamiento temporal
+        nlp_pipeline = load_pipeline_from_gcs(model_name, pipeline_task)
+        
+        # Realizar la predicción
         result = nlp_pipeline(input_text)
+        
         return {"response": result}
 
     except HTTPException as e:
@@ -113,10 +119,5 @@ def predict(
         raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
-    try:
-        import ssl
-        ssl_context = ssl.create_default_context()
-        uvicorn.run(app, host="0.0.0.0", port=7860)
-    except ModuleNotFoundError as e:
-        logger.error(f"Falta el módulo requerido: {str(e)}")
-        raise RuntimeError("El entorno no tiene instalado el paquete necesario para manejar conexiones seguras.")
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)