Update app.py

app.py

@@ -5,24 +5,18 @@ import boto3
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import JSONResponse
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from huggingface_hub import hf_hub_download
 import asyncio
+import concurrent.futures
 
-# Configuración del logger
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-console_handler = logging.StreamHandler()
-formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
-console_handler.setFormatter(formatter)
-logger.addHandler(console_handler)
 
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
 AWS_REGION = os.getenv("AWS_REGION")
 S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
-HUGGINGFACE_HUB_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
 
-MAX_TOKENS = 1024
+MAX_TOKENS = 1024  # Limite de tokens por fragmento
 
 s3_client = boto3.client(
     's3',
@@ -88,8 +82,9 @@ class S3DirectStream:
             model_files = await self.get_model_file_parts(model_prefix)
 
             if not model_files:
-                await self.download_and_upload_to_s3(model_prefix, model)
+                raise HTTPException(status_code=404, detail=f"Archivos del modelo {model_name} no encontrados en S3.")
 
+            # Verificar que existe el archivo config.json
             config_stream = await self.stream_from_s3(f"{model_prefix}/config.json")
             config_data = config_stream.read()
 
@@ -114,7 +109,7 @@ class S3DirectStream:
             tokenizer_stream = await self.stream_from_s3(f"{profile}/{model}/tokenizer.json")
             tokenizer_data = tokenizer_stream.read().decode("utf-8")
 
-            tokenizer = AutoTokenizer.from_pretrained(f"s3://{self.bucket_name}/{profile}/{model}")
+            tokenizer = AutoTokenizer.from_pretrained(f"{profile}/{model}")
             return tokenizer
         except Exception as e:
             raise HTTPException(status_code=500, detail=f"Error al cargar el tokenizer desde S3: {e}")
@@ -138,22 +133,6 @@ class S3DirectStream:
         except self.s3_client.exceptions.ClientError:
             return False
 
-    async def download_and_upload_to_s3(self, model_prefix, model_name):
-        try:
-            config_file = hf_hub_download(repo_id=model_name, filename="config.json", token=HUGGINGFACE_HUB_TOKEN)
-            tokenizer_file = hf_hub_download(repo_id=model_name, filename="tokenizer.json", token=HUGGINGFACE_HUB_TOKEN)
-
-            if not await self.file_exists_in_s3(f"{model_prefix}/config.json"):
-                with open(config_file, "rb") as file:
-                    self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_prefix}/config.json", Body=file)
-
-            if not await self.file_exists_in_s3(f"{model_prefix}/tokenizer.json"):
-                with open(tokenizer_file, "rb") as file:
-                    self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_prefix}/tokenizer.json", Body=file)
-
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al descargar o cargar archivos desde Hugging Face a S3: {e}")
-
 def split_text_by_tokens(text, tokenizer, max_tokens=MAX_TOKENS):
     tokens = tokenizer.encode(text)
     chunks = []
@@ -169,7 +148,7 @@ def continue_generation(input_text, model, tokenizer, max_tokens=MAX_TOKENS):
         input_text = tokenizer.decode(tokens[:max_tokens])
         output = model.generate(input_ids=tokenizer.encode(input_text, return_tensors="pt").input_ids)
         generated_text += tokenizer.decode(output[0], skip_special_tokens=True)
-        input_text = input_text[len(input_text):]
+        input_text = input_text[len(input_text):]  # Si la entrada se agotó, ya no hay más que procesar
     return generated_text
 
 @app.post("/predict/")
@@ -184,7 +163,7 @@ async def predict(model_request: dict):
 
         streamer = S3DirectStream(S3_BUCKET_NAME)
         
-        await streamer.create_s3_folders(model_name)
+        await streamer.create_s3_folders(model_name)  # Crear las carpetas si no existen
 
         model = await streamer.load_model_from_s3(model_name)
         tokenizer = await streamer.load_tokenizer_from_s3(model_name)
@@ -196,22 +175,18 @@ async def predict(model_request: dict):
 
         result = await asyncio.to_thread(nlp_pipeline, input_text)
 
-        if len(result) > MAX_TOKENS:
-            chunks = split_text_by_tokens(result, tokenizer)
+        chunks = split_text_by_tokens(result, tokenizer)
+
+        if len(chunks) > 1:
             full_result = ""
             for chunk in chunks:
                 full_result += continue_generation(chunk, model, tokenizer)
-            return {"result": full_result}
-        
-        return {"result": result}
-
-    except HTTPException as e:
-        logger.error(f"Error al realizar la predicción: {str(e.detail)}")
-        return JSONResponse(status_code=e.status_code, content={"detail": str(e.detail)})
+            return JSONResponse(content={"result": full_result})
+        else:
+            return JSONResponse(content={"result": result})
 
     except Exception as e:
-        logger.error(f"Error inesperado: {str(e)}")
-        return JSONResponse(status_code=500, content={"detail": "Error inesperado. Intenta más tarde."})
+        raise HTTPException(status_code=500, detail=f"Error al realizar la predicción: {e}")
 
 if __name__ == "__main__":
     import uvicorn