Update app.py

app.py

@@ -5,13 +5,16 @@ import logging
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from google.cloud import storage
+from google.auth import exceptions
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 import uvicorn
-import torch
+from google.cloud.storage.blob import Blob
 import requests
 import io
 from safetensors import safe_open
+import torch
 
+# Cargar las variables de entorno
 from dotenv import load_dotenv
 load_dotenv()
 
@@ -20,9 +23,11 @@ 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 del logger
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
+# Inicializar el cliente de Google Cloud Storage
 try:
     credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
     storage_client = storage.Client.from_service_account_info(credentials_info)
@@ -66,12 +71,13 @@ def load_model_from_gcs(model_name: str, model_files: list):
     gcs_handler = GCSHandler(GCS_BUCKET_NAME)
     model_blobs = {file: gcs_handler.download_file(f"{model_name}/{file}") for file in model_files}
     
+    # Verificar si el modelo es de safetensors o torch
     model_stream = model_blobs.get("pytorch_model.bin") or model_blobs.get("model.safetensors")
     config_stream = model_blobs.get("config.json")
     tokenizer_stream = model_blobs.get("tokenizer.json")
     
     if "safetensors" in model_stream.name:
-        model = load_safetensors_model(model_stream)
+        model = load_safetensors_model(model_stream, config_stream)
     else:
         model = AutoModelForCausalLM.from_pretrained(model_stream, config=config_stream)
     
@@ -79,7 +85,7 @@ def load_model_from_gcs(model_name: str, model_files: list):
     
     return model, tokenizer
 
-def load_safetensors_model(model_stream):
+def load_safetensors_model(model_stream, config_stream):
     with safe_open(model_stream, framework="pt") as model_data:
         model = torch.load(model_data)
     return model
@@ -88,7 +94,7 @@ def get_model_files_from_gcs(model_name: str):
     gcs_handler = GCSHandler(GCS_BUCKET_NAME)
     blob_list = list(gcs_handler.bucket.list_blobs(prefix=f"{model_name}/"))
     model_files = [blob.name for blob in blob_list if "pytorch_model" in blob.name or "model" in blob.name]
-    model_files = sorted(model_files)
+    model_files = sorted(model_files)  # Asegurar que los archivos fragmentados estén en el orden correcto
     return model_files
 
 @app.post("/predict/")
@@ -98,14 +104,17 @@ async def predict(request: DownloadModelRequest):
         gcs_handler = GCSHandler(GCS_BUCKET_NAME)
         model_prefix = request.model_name
         
+        # Obtener los archivos del modelo (incluyendo fragmentados)
         model_files = get_model_files_from_gcs(model_prefix)
         
         if not model_files:
             logger.error(f"Modelos no encontrados en GCS para '{model_prefix}'.")
             raise HTTPException(status_code=404, detail="Model files not found in GCS.")
 
+        # Cargar el modelo desde GCS
         model, tokenizer = load_model_from_gcs(model_prefix, model_files)
         
+        # Instanciar el pipeline de Hugging Face
         pipe = pipeline(request.pipeline_task, model=model, tokenizer=tokenizer)
 
         if request.pipeline_task in ["text-generation", "translation", "summarization"]: