Update app.py

app.py

@@ -6,7 +6,6 @@ from google.cloud import storage
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from pydantic import BaseModel
 from fastapi import FastAPI, HTTPException
-from io import BytesIO
 import requests
 import uvicorn
 from dotenv import load_dotenv
@@ -21,14 +20,9 @@ HF_API_TOKEN = os.getenv("HF_API_TOKEN")
 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)
-    logger.info(f"Conexión con Google Cloud Storage exitosa. Bucket: {GCS_BUCKET_NAME}")
-except Exception as e:
-    logger.error(f"Error al cargar las credenciales o bucket: {e}")
-    raise RuntimeError(f"Error al cargar las credenciales o bucket: {e}")
+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)
 
 app = FastAPI()
 
@@ -46,23 +40,16 @@ class GCSHandler:
 
     def create_folder_if_not_exists(self, folder_name):
         if not self.file_exists(folder_name):
-            logger.debug(f"Creando carpeta {folder_name} en GCS.")
             self.bucket.blob(folder_name + "/").upload_from_string("")
 
     def upload_file(self, blob_name, file_stream):
         self.create_folder_if_not_exists(os.path.dirname(blob_name))
         blob = self.bucket.blob(blob_name)
-        try:
-            blob.upload_from_file(file_stream)
-            logger.info(f"Archivo '{blob_name}' subido exitosamente a GCS.")
-        except Exception as e:
-            logger.error(f"Error subiendo el archivo '{blob_name}' a GCS: {e}")
-            raise HTTPException(status_code=500, detail=f"Error subiendo archivo '{blob_name}' a GCS")
+        blob.upload_from_file(file_stream)
 
     def download_file(self, blob_name):
         blob = self.bucket.blob(blob_name)
         if not blob.exists():
-            logger.error(f"Archivo '{blob_name}' no encontrado en GCS.")
             raise HTTPException(status_code=404, detail=f"File '{blob_name}' not found.")
         return blob.open("rb")
 
@@ -73,29 +60,21 @@ class GCSHandler:
 def download_model_from_huggingface(model_name):
     url = f"https://huggingface.co/{model_name}/tree/main"
     headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
-
-    try:
-        logger.info(f"Descargando el modelo '{model_name}' desde Hugging Face...")
-        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"lilmeaty_garca/{model_name}/{file_name}"
-                bucket.blob(blob_name).upload_from_string(file_content)
-                logger.info(f"Archivo '{file_name}' subido exitosamente al bucket GCS.")
-        else:
-            logger.error(f"Error al acceder al árbol de archivos de Hugging Face para '{model_name}'.")
-            raise HTTPException(status_code=404, detail="Error al acceder al árbol de archivos de Hugging Face.")
-    except Exception as e:
-        logger.error(f"Error descargando archivos de Hugging Face: {e}")
-        raise HTTPException(status_code=500, detail=f"Error descargando archivos de Hugging Face: {e}")
+    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"models/{model_name}/{file_name}"
+            bucket.blob(blob_name).upload_from_string(file_content)
+    else:
+        raise HTTPException(status_code=404, detail="Error accessing Hugging Face model files.")
 
 def download_and_verify_model(model_name):
     model_files = [
@@ -105,8 +84,7 @@ def download_and_verify_model(model_name):
         "model.safetensors",
     ]
     gcs_handler = GCSHandler(GCS_BUCKET_NAME)
-    model_files_exist = all(gcs_handler.file_exists(f"lilmeaty_garca/{model_name}/{file}") for file in model_files)
-    if not model_files_exist:
+    if not all(gcs_handler.file_exists(f"models/{model_name}/{file}") for file in model_files):
         download_model_from_huggingface(model_name)
 
 def load_model_from_gcs(model_name):
@@ -117,103 +95,58 @@ def load_model_from_gcs(model_name):
         "model.safetensors",
     ]
     gcs_handler = GCSHandler(GCS_BUCKET_NAME)
-    model_files_streams = {}
-
-    for file in model_files:
-        file_path = f"lilmeaty_garca/{model_name}/{file}"
-        if gcs_handler.file_exists(file_path):
-            model_files_streams[file] = gcs_handler.download_file(file_path)
-        else:
-            logger.error(f"Archivo '{file}' no encontrado en GCS.")
-            raise HTTPException(status_code=500, detail=f"Archivo '{file}' no encontrado.")
-
+    model_files_streams = {
+        file: gcs_handler.download_file(f"models/{model_name}/{file}")
+        for file in model_files if gcs_handler.file_exists(f"models/{model_name}/{file}")
+    }
     model_stream = model_files_streams.get("pytorch_model.bin") or model_files_streams.get("model.safetensors")
     tokenizer_stream = model_files_streams.get("tokenizer.json")
     config_stream = model_files_streams.get("config.json")
-
     model = AutoModelForCausalLM.from_pretrained(model_stream, config=config_stream)
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
-
     return model, tokenizer
 
 def load_model(model_name):
     gcs_handler = GCSHandler(GCS_BUCKET_NAME)
     try:
-        model, tokenizer = load_model_from_gcs(model_name)
-        logger.info(f"Modelo '{model_name}' cargado exitosamente desde GCS.")
-        return model, tokenizer
+        return load_model_from_gcs(model_name)
     except HTTPException:
-        logger.warning(f"Modelo '{model_name}' no encontrado en GCS. Intentando como model_id...")
-
-    try:
         download_and_verify_model(model_name)
-        model, tokenizer = load_model_from_gcs(model_name)
-        logger.info(f"Modelo '{model_name}' cargado exitosamente desde Hugging Face.")
-        return model, tokenizer
-    except Exception as e:
-        logger.error(f"Error al intentar cargar el modelo '{model_name}': {e}")
-        raise HTTPException(status_code=500, detail=f"Error al cargar el modelo '{model_name}': {e}")
+        return load_model_from_gcs(model_name)
 
 @app.on_event("startup")
 async def startup():
-    try:
-        logger.info("Iniciando la descarga de modelos en segundo plano...")
-        run_in_background()
-        gcs_handler = GCSHandler(GCS_BUCKET_NAME)
-        blobs = list(bucket.list_blobs(prefix="lilmeaty_garca/"))
-        model_names = set([blob.name.split("/")[1] for blob in blobs])
-
-        def download_model_thread(model_name):
-            try:
-                download_and_verify_model(model_name)
-            except Exception as e:
-                logger.error(f"Error descargando modelo '{model_name}': {e}")
-
-        threads = []
-        for model_name in model_names:
-            thread = threading.Thread(target=download_model_thread, args=(model_name,))
-            thread.start()
-            threads.append(thread)
-
-        for thread in threads:
-            thread.join()
-
-        logger.info("Todos los modelos se descargaron correctamente o ya estaban presentes.")
-    except Exception as e:
-        logger.error(f"Error durante la descarga de modelos al iniciar: {e}")
-        raise HTTPException(status_code=500, detail=f"Error durante la descarga de modelos: {e}")
+    gcs_handler = GCSHandler(GCS_BUCKET_NAME)
+    blobs = list(bucket.list_blobs(prefix="models/"))
+    model_names = set(blob.name.split("/")[1] for blob in blobs)
+    def download_model_thread(model_name):
+        try:
+            download_and_verify_model(model_name)
+        except Exception as e:
+            logger.error(f"Error downloading model '{model_name}': {e}")
+    threads = [threading.Thread(target=download_model_thread, args=(model_name,)) for model_name in model_names]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()
 
 @app.post("/predict/")
 async def predict(request: DownloadModelRequest):
-    try:
-        model_name = request.model_name
-        pipeline_task = request.pipeline_task
-        input_text = request.input_text
-
-        model, tokenizer = load_model(model_name)
-        pipe = pipeline(pipeline_task, model=model, tokenizer=tokenizer)
-        result = pipe(input_text)
-
-        return {"result": result}
-    except Exception as e:
-        logger.error(f"Error procesando la solicitud: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+    model_name = request.model_name
+    pipeline_task = request.pipeline_task
+    input_text = request.input_text
+    model, tokenizer = load_model(model_name)
+    pipe = pipeline(pipeline_task, model=model, tokenizer=tokenizer)
+    result = pipe(input_text)
+    return {"result": result}
 
 def download_all_models_in_background():
     models_url = "https://huggingface.co/api/models"
-    try:
-        response = requests.get(models_url)
-        if response.status_code != 200:
-            logger.error("Error al obtener la lista de modelos de Hugging Face.")
-            raise HTTPException(status_code=500, detail="Error al obtener la lista de modelos.")
-
+    response = requests.get(models_url)
+    if response.status_code == 200:
         models = response.json()
         for model in models:
-            model_name = model["id"]
-            download_model_from_huggingface(model_name)
-    except Exception as e:
-        logger.error(f"Error al descargar modelos en segundo plano: {e}")
-        raise HTTPException(status_code=500, detail="Error al descargar modelos en segundo plano.")
+            download_model_from_huggingface(model["id"])
 
 def run_in_background():
     threading.Thread(target=download_all_models_in_background, daemon=True).start()