Update app.py

app.py

@@ -1,8 +1,8 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
 import os
 import json
 import requests
+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
@@ -49,28 +49,24 @@ class GCSHandler:
         return BytesIO(blob.download_as_bytes())
 
 def download_model_from_huggingface(model_name):
-    file_patterns = [
-        "pytorch_model.bin",
-        "config.json",
-        "tokenizer.json",
-        "model.safetensors",
-    ]
-    for i in range(1, 100):
-        file_patterns.extend([f"pytorch_model-{i:05}-of-00001", f"model-{i:05}"])
+    url = f"https://huggingface.co/{model_name}/tree/main"
+    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
     
-    # Descargar los archivos del modelo
-    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:
+    # Intentar obtener el árbol de archivos
+    try:
+        response = requests.get(url, headers=headers)
+        if response.status_code == 200:
+            # Extraer la lista de archivos del árbol (parseo HTML o JSON depende de la respuesta)
+            # Aquí asumimos que el archivo de modelos está disponible
+            file_urls = []  # Aquí agregarías la lógica para extraer los enlaces correctos del HTML de la página
+            for file_url in file_urls:
+                filename = file_url.split("/")[-1]
                 blob_name = f"{model_name}/{filename}"
-                bucket.blob(blob_name).upload_from_file(BytesIO(response.content))
-            else:
-                raise HTTPException(status_code=404, detail=f"File {filename} not found on Hugging Face.")
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error downloading {filename} from Hugging Face: {e}")
+                bucket.blob(blob_name).upload_from_file(BytesIO(requests.get(file_url).content))
+        else:
+            raise HTTPException(status_code=404, detail="Error al acceder al árbol de archivos de Hugging Face.")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error descargando archivos de Hugging Face: {e}")
 
 @app.post("/predict/")
 async def predict(request: DownloadModelRequest):
@@ -83,8 +79,6 @@ async def predict(request: DownloadModelRequest):
             "tokenizer.json",
             "model.safetensors",
         ]
-        for i in range(1, 100):
-            model_files.extend([f"pytorch_model-{i:05}-of-00001", f"model-{i:05}"])
         
         # Verificar si los archivos del modelo están en GCS
         model_files_exist = all(gcs_handler.file_exists(f"{model_prefix}/{file}") for file in model_files)