Update app.py

app.py

@@ -56,7 +56,6 @@ def download_model_from_huggingface(model_name):
     try:
         response = requests.get(url, headers=headers)
         if response.status_code == 200:
-            # Enlace a los archivos del modelo
             model_files = [
                 "pytorch_model.bin",
                 "config.json",
@@ -85,17 +84,13 @@ async def predict(request: DownloadModelRequest):
             "model.safetensors",
         ]
         
-        # 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")
         model_stream = model_files_streams.get("pytorch_model.bin")
@@ -103,7 +98,6 @@ async def predict(request: DownloadModelRequest):
         if not config_stream or not tokenizer_stream or not model_stream:
             raise HTTPException(status_code=500, detail="Required model files missing.")
         
-        # 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")
@@ -116,14 +110,11 @@ async def predict(request: DownloadModelRequest):
             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(tmp_dir)
+            model = AutoModelForCausalLM.from_pretrained(tmp_dir, from_tf=False)
             tokenizer = AutoTokenizer.from_pretrained(tmp_dir)
             
-            # 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}