Update app.py

app.py

@@ -57,6 +57,8 @@ def download_model_from_huggingface(model_name):
     ]
     for i in range(1, 100):
         file_patterns.extend([f"pytorch_model-{i:05}-of-00001", f"model-{i:05}"])
+    
+    # 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}"}
@@ -65,6 +67,8 @@ def download_model_from_huggingface(model_name):
             if response.status_code == 200:
                 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}")
 
@@ -81,18 +85,38 @@ async def predict(request: DownloadModelRequest):
         ]
         for i in range(1, 100):
             model_files.extend([f"pytorch_model-{i:05}-of-00001", f"model-{i:05}"])
-        if not any(gcs_handler.file_exists(f"{model_prefix}/{file}") for file in model_files):
+        
+        # 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)
+        
+        if not model_files_exist:
+            # Descargar el modelo si no existe
             download_model_from_huggingface(model_prefix)
+        
+        # Descargar los archivos necesarios
         model_files_streams = {file: gcs_handler.download_file(f"{model_prefix}/{file}") for file in model_files if gcs_handler.file_exists(f"{model_prefix}/{file}")}
+        
+        # Asegurar que los archivos esenciales estén presentes
         config_stream = model_files_streams.get("config.json")
         tokenizer_stream = model_files_streams.get("tokenizer.json")
+        
         if not config_stream or not tokenizer_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)
+        
         return {"response": result}
+    
+    except HTTPException as e:
+        raise e
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error: {e}")