Update app.py

app.py

@@ -9,6 +9,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from io import BytesIO
 from dotenv import load_dotenv
 import uvicorn
+import tempfile
 
 load_dotenv()
 
@@ -52,17 +53,21 @@ def download_model_from_huggingface(model_name):
     url = f"https://huggingface.co/{model_name}/tree/main"
     headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
     
-    # 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(requests.get(file_url).content))
+            # Enlace a los archivos del modelo
+            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"{model_name}/{file_name}"
+                bucket.blob(blob_name).upload_from_file(BytesIO(file_content))
         else:
             raise HTTPException(status_code=404, detail="Error al acceder al árbol de archivos de Hugging Face.")
     except Exception as e:
@@ -93,19 +98,33 @@ async def predict(request: DownloadModelRequest):
         # Asegurar que los archivos esenciales estén presentes
         config_stream = model_files_streams.get("config.json")
         tokenizer_stream = model_files_streams.get("tokenizer.json")
+        model_stream = model_files_streams.get("pytorch_model.bin")
         
-        if not config_stream or not tokenizer_stream:
+        if not config_stream or not tokenizer_stream or not model_stream:
             raise HTTPException(status_code=500, detail="Required model files missing.")
         
-        # Cargar el modelo y el tokenizador
-        model = AutoModelForCausalLM.from_pretrained(config_stream)
-        tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
-        
-        # Crear un pipeline para la tarea deseada
-        pipeline_ = pipeline(request.pipeline_task, model=model, tokenizer=tokenizer)
-        
-        # Realizar la predicción
-        result = pipeline_(request.input_text)
+        # Guardar los archivos en directorios temporales
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            config_path = os.path.join(tmp_dir, "config.json")
+            tokenizer_path = os.path.join(tmp_dir, "tokenizer.json")
+            model_path = os.path.join(tmp_dir, "pytorch_model.bin")
+
+            with open(config_path, 'wb') as f:
+                f.write(config_stream.read())
+            with open(tokenizer_path, 'wb') as f:
+                f.write(tokenizer_stream.read())
+            with open(model_path, 'wb') as f:
+                f.write(model_stream.read())
+            
+            # Cargar el modelo y el tokenizador desde los archivos temporales
+            model = AutoModelForCausalLM.from_pretrained(model_path)
+            tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+            
+            # Crear un pipeline para la tarea deseada
+            pipeline_ = pipeline(request.pipeline_task, model=model, tokenizer=tokenizer)
+            
+            # Realizar la predicción
+            result = pipeline_(request.input_text)
         
         return {"response": result}