Update app.py

app.py

@@ -7,8 +7,7 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from google.cloud import storage
 from google.auth import exceptions
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from transformers.hf_api import HfApi, HfFolder, HfLoginManager
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from io import BytesIO
 from dotenv import load_dotenv
 import uvicorn
@@ -33,16 +32,6 @@ except (exceptions.DefaultCredentialsError, json.JSONDecodeError, KeyError) as e
 # Inicialización de FastAPI
 app = FastAPI()
 
-# Inicio de sesión en Hugging Face
-try:
-    if not HF_API_TOKEN:
-        raise ValueError("El token de Hugging Face no está definido en las variables de entorno.")
-    HfApi().set_access_token(HF_API_TOKEN)
-    print("Inicio de sesión en Hugging Face exitoso.")
-except Exception as e:
-    print(f"Error al iniciar sesión en Hugging Face: {e}")
-    exit(1)
-
 
 class DownloadModelRequest(BaseModel):
     model_name: str
@@ -88,28 +77,65 @@ class GCSStreamHandler:
         return model_files
 
 
+def download_model_from_huggingface(model_name):
+    """
+    Descarga un modelo desde Hugging Face y lo sube a GCS en streaming.
+    """
+    file_patterns = [
+        "pytorch_model.bin",
+        "model.safetensors",
+        "config.json",
+        "tokenizer.json",
+    ]
+
+    # Agregar patrones para fragmentos de modelos
+    for i in range(1, 100):
+        file_patterns.append(f"pytorch_model-{i:05}-of-{100:05}")
+        file_patterns.append(f"model-{i:05}")
+
+    for filename in file_patterns:
+        url = f"https://huggingface.co/{model_name}/resolve/main/{filename}"
+        headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
+        try:
+            response = requests.get(url, headers=headers, stream=True)
+            if response.status_code == 200:
+                blob_name = f"{model_name}/{filename}"
+                blob = bucket.blob(blob_name)
+                blob.upload_from_file(BytesIO(response.content))
+                print(f"Archivo {filename} subido correctamente a GCS.")
+        except Exception as e:
+            print(f"Archivo {filename} no encontrado: {e}")
+
+
 @app.post("/predict/")
 async def predict(request: DownloadModelRequest):
+    """
+    Endpoint para realizar predicciones. Si el modelo no existe en GCS, se descarga automáticamente.
+    """
     try:
         gcs_handler = GCSStreamHandler(GCS_BUCKET_NAME)
 
-        # Asegura la estructura del bucket
-        gcs_handler.ensure_bucket_structure(request.model_name)
-
-        # Define patrones para los archivos de modelos
+        # Verificar si el modelo ya está en GCS
+        model_prefix = request.model_name
         model_patterns = [
             r"pytorch_model-\d+-of-\d+",
             r"model-\d+",
             r"pytorch_model.bin",
-            r"model.safetensors"
+            r"model.safetensors",
         ]
 
-        # Carga los archivos del modelo desde el bucket
-        model_files = gcs_handler.stream_model_files(request.model_name, model_patterns)
+        if not any(
+            gcs_handler.file_exists(f"{model_prefix}/{pattern}") for pattern in model_patterns
+        ):
+            print(f"Modelo {model_prefix} no encontrado en GCS. Descargando desde Hugging Face...")
+            download_model_from_huggingface(model_prefix)
+
+        # Carga archivos del modelo desde GCS
+        model_files = gcs_handler.stream_model_files(model_prefix, model_patterns)
 
-        # Cargar configuración y modelo
-        config_stream = gcs_handler.stream_file_from_gcs(f"{request.model_name}/config.json")
-        tokenizer_stream = gcs_handler.stream_file_from_gcs(f"{request.model_name}/tokenizer.json")
+        # Configuración y tokenización
+        config_stream = gcs_handler.stream_file_from_gcs(f"{model_prefix}/config.json")
+        tokenizer_stream = gcs_handler.stream_file_from_gcs(f"{model_prefix}/tokenizer.json")
 
         model = AutoModelForCausalLM.from_pretrained(BytesIO(config_stream))
         state_dict = {}
@@ -135,42 +161,5 @@ async def predict(request: DownloadModelRequest):
         raise HTTPException(status_code=500, detail=f"Error: {e}")
 
 
-@app.post("/upload/")
-async def upload_model_to_gcs(model_name: str):
-    """
-    Descarga un modelo desde Hugging Face y lo sube a GCS en streaming.
-    """
-    try:
-        gcs_handler = GCSStreamHandler(GCS_BUCKET_NAME)
-
-        # Archivos comunes de los modelos
-        file_patterns = [
-            "pytorch_model.bin",
-            "model.safetensors",
-            "config.json",
-            "tokenizer.json",
-        ]
-
-        # Agregar patrones para fragmentos de modelos
-        for i in range(1, 100):
-            file_patterns.append(f"pytorch_model-{i:05}-of-{100:05}")
-            file_patterns.append(f"model-{i:05}")
-
-        for filename in file_patterns:
-            url = f"https://huggingface.co/{model_name}/resolve/main/{filename}"
-            headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
-            try:
-                response = requests.get(url, headers=headers, stream=True)
-                if response.status_code == 200:
-                    blob_name = f"{model_name}/{filename}"
-                    blob = bucket.blob(blob_name)
-                    blob.upload_from_file(BytesIO(response.content))
-                    print(f"Archivo {filename} subido correctamente a GCS.")
-            except Exception as e:
-                print(f"Archivo {filename} no encontrado: {e}")
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error al subir modelo: {e}")
-
-
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8000)