Update app.py

app.py

@@ -7,7 +7,6 @@ from fastapi.responses import JSONResponse
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from huggingface_hub import hf_hub_download
 from tqdm import tqdm
-import io
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -50,25 +49,20 @@ class S3DirectStream:
 
     def stream_from_s3(self, key):
         try:
-            logger.info(f"Descargando {key} desde S3...")
             response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
             return response['Body']
         except self.s3_client.exceptions.NoSuchKey:
-            logger.error(f"El archivo {key} no existe en el bucket S3.")
             raise HTTPException(status_code=404, detail=f"El archivo {key} no existe en el bucket S3.")
         except Exception as e:
-            logger.error(f"Error al descargar {key} desde S3: {str(e)}")
             raise HTTPException(status_code=500, detail=f"Error al descargar {key} desde S3: {str(e)}")
 
     def get_model_file_parts(self, model_name):
         try:
             model_prefix = model_name.lower()
-            logger.info(f"Obteniendo archivos para el modelo {model_name} desde S3...")
             files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_prefix)
             model_files = [obj['Key'] for obj in files.get('Contents', []) if model_prefix in obj['Key']]
             return model_files
         except Exception as e:
-            logger.error(f"Error al obtener archivos del modelo {model_name} desde S3: {e}")
             raise HTTPException(status_code=500, detail=f"Error al obtener archivos del modelo {model_name} desde S3: {e}")
 
     def load_model_from_s3(self, model_name):
@@ -77,20 +71,16 @@ class S3DirectStream:
             model_files = self.get_model_file_parts(model_prefix)
 
             if not model_files:
-                logger.info(f"El modelo {model_name} no está en S3, descargando desde Hugging Face...")
                 self.download_and_upload_from_huggingface(model_name)
                 model_files = self.get_model_file_parts(model_prefix)
 
             if not model_files:
-                logger.error(f"Archivos del modelo {model_name} no encontrados en S3.")
                 raise HTTPException(status_code=404, detail=f"Archivos del modelo {model_name} no encontrados en S3.")
 
-            logger.info(f"Cargando archivos del modelo {model_name}...")
             config_stream = self.stream_from_s3(f"{model_prefix}/config.json")
             config_data = config_stream.read()
 
             if not config_data:
-                logger.error(f"El archivo de configuración {model_prefix}/config.json está vacío.")
                 raise HTTPException(status_code=500, detail=f"El archivo de configuración {model_prefix}/config.json está vacío.")
             
             config_text = config_data.decode("utf-8")
@@ -100,23 +90,19 @@ class S3DirectStream:
             return model
 
         except Exception as e:
-            logger.error(f"Error al cargar el modelo desde S3: {e}")
             raise HTTPException(status_code=500, detail=f"Error al cargar el modelo desde S3: {e}")
 
     def load_tokenizer_from_s3(self, model_name):
         try:
-            logger.info(f"Cargando el tokenizer del modelo {model_name} desde S3...")
             tokenizer_stream = self.stream_from_s3(f"{model_name}/tokenizer.json")
             tokenizer_data = tokenizer_stream.read().decode("utf-8")
             tokenizer = AutoTokenizer.from_pretrained(f"s3://{self.bucket_name}/{model_name}")
             return tokenizer
         except Exception as e:
-            logger.error(f"Error al cargar el tokenizer desde S3: {e}")
             raise HTTPException(status_code=500, detail=f"Error al cargar el tokenizer desde S3: {e}")
 
     def download_and_upload_from_huggingface(self, model_name):
         try:
-            logger.info(f"Descargando modelo {model_name} desde Hugging Face...")
             files_to_download = hf_hub_download(repo_id=model_name, use_auth_token=HUGGINGFACE_TOKEN, local_dir=model_name)
 
             for file in tqdm(files_to_download, desc="Subiendo archivos a S3"):
@@ -126,7 +112,6 @@ class S3DirectStream:
                     self.upload_file_to_s3(file, s3_key)
 
         except Exception as e:
-            logger.error(f"Error al descargar y subir modelo desde Hugging Face: {e}")
             raise HTTPException(status_code=500, detail=f"Error al descargar y subir modelo desde Hugging Face: {e}")
 
     def upload_file_to_s3(self, file_path, s3_key):
@@ -134,9 +119,7 @@ class S3DirectStream:
             with open(file_path, 'rb') as data:
                 self.s3_client.put_object(Bucket=self.bucket_name, Key=s3_key, Body=data)
             os.remove(file_path)
-            logger.info(f"Archivo {file_path} subido correctamente a S3 y eliminado localmente.")
         except Exception as e:
-            logger.error(f"Error al subir archivo a S3: {e}")
             raise HTTPException(status_code=500, detail=f"Error al subir archivo a S3: {e}")
 
     def file_exists_in_s3(self, s3_key):
@@ -161,7 +144,6 @@ async def predict(model_request: dict):
         tokenizer = streamer.load_tokenizer_from_s3(model_name)
 
         if task not in PIPELINE_MAP:
-            logger.error("Pipeline task no soportado")
             raise HTTPException(status_code=400, detail="Pipeline task no soportado")
 
         nlp_pipeline = pipeline(PIPELINE_MAP[task], model=model, tokenizer=tokenizer)
@@ -174,7 +156,6 @@ async def predict(model_request: dict):
             return JSONResponse(content={"result": result})
 
     except Exception as e:
-        logger.error(f"Error al realizar la predicción: {e}")
         raise HTTPException(status_code=500, detail=f"Error al realizar la predicción: {e}")
 
 if __name__ == "__main__":