Update app.py

app.py

@@ -5,8 +5,8 @@ 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
-from tqdm import tqdm
+import asyncio
+import concurrent.futures
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -15,7 +15,8 @@ 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_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+
+MAX_TOKENS = 1024  # Limite de tokens por fragmento
 
 s3_client = boto3.client(
     's3',
@@ -47,7 +48,11 @@ class S3DirectStream:
         )
         self.bucket_name = bucket_name
 
-    def stream_from_s3(self, key):
+    async def stream_from_s3(self, key):
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(None, self._stream_from_s3, key)
+
+    def _stream_from_s3(self, key):
         try:
             response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
             return response['Body']
@@ -56,7 +61,11 @@ class S3DirectStream:
         except Exception as e:
             raise HTTPException(status_code=500, detail=f"Error al descargar {key} desde S3: {str(e)}")
 
-    def get_model_file_parts(self, model_name):
+    async def get_model_file_parts(self, model_name):
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(None, self._get_model_file_parts, model_name)
+
+    def _get_model_file_parts(self, model_name):
         try:
             model_prefix = model_name.lower()
             files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_prefix)
@@ -65,21 +74,17 @@ class S3DirectStream:
         except Exception as 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):
+    async def load_model_from_s3(self, model_name):
         try:
             profile, model = model_name.split("/", 1) if "/" in model_name else ("", model_name)
 
             model_prefix = f"{profile}/{model}".lower()
-            model_files = self.get_model_file_parts(model_prefix)
-
-            if not model_files:
-                self.download_and_upload_from_huggingface(model_name)
-                model_files = self.get_model_file_parts(model_prefix)
+            model_files = await self.get_model_file_parts(model_prefix)
 
             if not model_files:
                 raise HTTPException(status_code=404, detail=f"Archivos del modelo {model_name} no encontrados en S3.")
 
-            config_stream = self.stream_from_s3(f"{model_prefix}/config.json")
+            config_stream = await self.stream_from_s3(f"{model_prefix}/config.json")
             config_data = config_stream.read()
 
             if not config_data:
@@ -94,18 +99,13 @@ class S3DirectStream:
         except HTTPException as e:
             raise e
         except Exception as e:
-            try:
-                logger.error(f"Error al cargar el modelo desde S3, intentando desde Hugging Face: {e}")
-                model = AutoModelForCausalLM.from_pretrained(model_name)
-                return model
-            except Exception as hf_error:
-                raise HTTPException(status_code=500, detail=f"Error al cargar el modelo desde Hugging Face: {hf_error}")
-
-    def load_tokenizer_from_s3(self, model_name):
+            raise HTTPException(status_code=500, detail=f"Error al cargar el modelo desde S3: {e}")
+
+    async def load_tokenizer_from_s3(self, model_name):
         try:
             profile, model = model_name.split("/", 1) if "/" in model_name else ("", model_name)
 
-            tokenizer_stream = self.stream_from_s3(f"{profile}/{model}/tokenizer.json")
+            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"{profile}/{model}")
@@ -113,46 +113,42 @@ class S3DirectStream:
         except Exception as 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:
-            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"):
-                file_name = os.path.basename(file)
-                profile, model = model_name.split("/", 1) if "/" in model_name else ("", model_name)
-                s3_key = f"{profile}/{model}/{file_name}"
-                if not self.file_exists_in_s3(s3_key):
-                    self.upload_file_to_s3(file, s3_key)
-
-        except Exception as 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):
-        try:
-            self.create_s3_folders(s3_key)
-            s3_client.put_object(Bucket=self.bucket_name, Key=s3_key, Body=open(file_path, 'rb'))
-            os.remove(file_path)
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al subir archivo a S3: {e}")
-
-    def create_s3_folders(self, s3_key):
+    async def create_s3_folders(self, s3_key):
         try:
             folder_keys = s3_key.split('/')
             for i in range(1, len(folder_keys)):
                 folder_key = '/'.join(folder_keys[:i]) + '/'
-                if not self.file_exists_in_s3(folder_key):
+                if not await self.file_exists_in_s3(folder_key):
                     self.s3_client.put_object(Bucket=self.bucket_name, Key=folder_key, Body='')
 
         except Exception as e:
             raise HTTPException(status_code=500, detail=f"Error al crear carpetas en S3: {e}")
 
-    def file_exists_in_s3(self, s3_key):
+    async def file_exists_in_s3(self, s3_key):
         try:
             self.s3_client.head_object(Bucket=self.bucket_name, Key=s3_key)
             return True
         except self.s3_client.exceptions.ClientError:
             return False
 
+def split_text_by_tokens(text, tokenizer, max_tokens=MAX_TOKENS):
+    tokens = tokenizer.encode(text)
+    chunks = []
+    for i in range(0, len(tokens), max_tokens):
+        chunk = tokens[i:i+max_tokens]
+        chunks.append(tokenizer.decode(chunk))
+    return chunks
+
+def continue_generation(input_text, model, tokenizer, max_tokens=MAX_TOKENS):
+    generated_text = ""
+    while len(input_text) > 0:
+        tokens = tokenizer.encode(input_text)
+        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):]  # Si la entrada se agotó, ya no hay más que procesar
+    return generated_text
+
 @app.post("/predict/")
 async def predict(model_request: dict):
     try:
@@ -164,18 +160,23 @@ async def predict(model_request: dict):
             raise HTTPException(status_code=400, detail="Faltan parámetros en la solicitud.")
 
         streamer = S3DirectStream(S3_BUCKET_NAME)
-        model = streamer.load_model_from_s3(model_name)
-        tokenizer = streamer.load_tokenizer_from_s3(model_name)
+        model = await streamer.load_model_from_s3(model_name)
+        tokenizer = await streamer.load_tokenizer_from_s3(model_name)
 
         if task not in PIPELINE_MAP:
             raise HTTPException(status_code=400, detail="Pipeline task no soportado")
 
         nlp_pipeline = pipeline(PIPELINE_MAP[task], model=model, tokenizer=tokenizer)
 
-        result = nlp_pipeline(input_text)
+        result = await asyncio.to_thread(nlp_pipeline, input_text)
+
+        chunks = split_text_by_tokens(result, tokenizer)
 
-        if isinstance(result, dict) and 'file' in result:
-            return JSONResponse(content={"file": result['file']})
+        if len(chunks) > 1:
+            full_result = ""
+            for chunk in chunks:
+                full_result += continue_generation(chunk, model, tokenizer)
+            return JSONResponse(content={"result": full_result})
         else:
             return JSONResponse(content={"result": result})
 
@@ -184,4 +185,4 @@ async def predict(model_request: dict):
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=8000)