Update app.py

app.py

@@ -3,10 +3,9 @@ import json
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from huggingface_hub import hf_hub_download
 import boto3
 import logging
-import asyncio
+from huggingface_hub import hf_hub_download
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -21,8 +20,6 @@ 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
-
 s3_client = boto3.client(
     's3',
     aws_access_key_id=AWS_ACCESS_KEY_ID,
@@ -47,31 +44,33 @@ class S3DirectStream:
         )
         self.bucket_name = bucket_name
 
-    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):
+    async def download_and_upload_to_s3(self, model_name):
         try:
-            response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
-            return response['Body'].read()
-        except self.s3_client.exceptions.NoSuchKey:
-            raise HTTPException(status_code=404, detail=f"El archivo {key} no existe en el bucket S3.")
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al descargar {key} desde S3: {str(e)}")
+            model_name = model_name.replace("/", "-").lower()
 
-    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)
+            # Descarga de los archivos desde Hugging Face
+            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)
+
+            # Verifica si ya existen en S3
+            if not await self.file_exists_in_s3(f"{model_name}/config.json"):
+                with open(config_file, "rb") as file:
+                    self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/config.json", Body=file)
+
+            if not await self.file_exists_in_s3(f"{model_name}/tokenizer.json"):
+                with open(tokenizer_file, "rb") as file:
+                    self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/tokenizer.json", Body=file)
 
-    def _get_model_file_parts(self, model_name):
-        try:
-            model_name = model_name.replace("/", "-").lower()
-            files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_name)
-            model_files = [obj['Key'] for obj in files.get('Contents', []) if model_name in obj['Key']]
-            return model_files
         except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al obtener archivos del modelo {model_name} desde S3: {e}")
+            logger.error(f"Error al cargar el modelo desde Hugging Face a S3: {e}")
+            raise HTTPException(status_code=500, detail=f"Error al cargar el modelo: {str(e)}")
+
+    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
 
     async def load_model_from_s3(self, model_name):
         try:
@@ -81,97 +80,53 @@ class S3DirectStream:
             if not model_files:
                 await self.download_and_upload_to_s3(model_name)
 
+            # Cargar configuración del modelo
             config_data = await self.stream_from_s3(f"{model_name}/config.json")
-            if not config_data:
-                raise HTTPException(status_code=500, detail=f"El archivo de configuración {model_name}/config.json está vacío o no se pudo leer.")
-            
             if isinstance(config_data, bytes):
                 config_data = config_data.decode("utf-8")
             
             config_json = json.loads(config_data)
 
-            model = AutoModelForCausalLM.from_pretrained(f"s3://{self.bucket_name}/{model_name}", config=config_json, from_tf=False)
+            model = AutoModelForCausalLM.from_pretrained(f"s3://{self.bucket_name}/{model_name}", config=config_json)
             return model
 
         except HTTPException as e:
             raise e
         except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al cargar el modelo desde S3: {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: {str(e)}")
 
     async def load_tokenizer_from_s3(self, model_name):
         try:
             model_name = model_name.replace("/", "-").lower()
             tokenizer_data = await self.stream_from_s3(f"{model_name}/tokenizer.json")
-            
+
             if isinstance(tokenizer_data, bytes):
                 tokenizer_data = tokenizer_data.decode("utf-8") 
             
             tokenizer = AutoTokenizer.from_pretrained(f"s3://{self.bucket_name}/{model_name}")
             return tokenizer
         except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al cargar el tokenizer desde S3: {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: {str(e)}")
 
-    async def create_s3_folders(self, s3_key):
+    async def stream_from_s3(self, key):
         try:
-            folder_keys = s3_key.split('-')
-            for i in range(1, len(folder_keys)):
-                folder_key = '-'.join(folder_keys[:i]) + '/'
-                if not await self.file_exists_in_s3(folder_key):
-                    logger.info(f"Creando carpeta en S3: {folder_key}")
-                    self.s3_client.put_object(Bucket=self.bucket_name, Key=folder_key, Body='')
-
+            response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
+            return response['Body'].read()
+        except self.s3_client.exceptions.NoSuchKey:
+            raise HTTPException(status_code=404, detail=f"El archivo {key} no existe en el bucket S3.")
         except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al crear carpetas en S3: {e}")
-
-    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
+            raise HTTPException(status_code=500, detail=f"Error al descargar {key} desde S3: {str(e)}")
 
-    async def download_and_upload_to_s3(self, model_name, force_download=False):
+    async def get_model_file_parts(self, model_name):
         try:
-            if force_download:
-                logger.info(f"Forzando la descarga del modelo {model_name} y la carga a S3.")
-            
             model_name = model_name.replace("/", "-").lower()
-
-            if not await self.file_exists_in_s3(f"{model_name}/config.json") or not await self.file_exists_in_s3(f"{model_name}/tokenizer.json"):
-                config_file = hf_hub_download(repo_id=model_name, filename="config.json", token=HUGGINGFACE_HUB_TOKEN, force_download=force_download)
-                tokenizer_file = hf_hub_download(repo_id=model_name, filename="tokenizer.json", token=HUGGINGFACE_HUB_TOKEN, force_download=force_download)
-
-                await self.create_s3_folders(f"{model_name}/")
-
-                if not await self.file_exists_in_s3(f"{model_name}/config.json"):
-                    with open(config_file, "rb") as file:
-                        self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/config.json", Body=file)
-
-                if not await self.file_exists_in_s3(f"{model_name}/tokenizer.json"):
-                    with open(tokenizer_file, "rb") as file:
-                        self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/tokenizer.json", Body=file)
-            else:
-                logger.info(f"Los archivos del modelo {model_name} ya existen en S3. No es necesario descargarlos de nuevo.")
-
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al descargar o cargar archivos desde Hugging Face a S3: {e}")
-
-    async def resume_download(self, model_name):
-        try:
-            logger.info(f"Reanudando la descarga del modelo {model_name} desde Hugging Face.")
-            config_file = hf_hub_download(repo_id=model_name, filename="config.json", token=HUGGINGFACE_HUB_TOKEN, resume_download=True)
-            tokenizer_file = hf_hub_download(repo_id=model_name, filename="tokenizer.json", token=HUGGINGFACE_HUB_TOKEN, resume_download=True)
-
-            if not await self.file_exists_in_s3(f"{model_name}/config.json"):
-                with open(config_file, "rb") as file:
-                    self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/config.json", Body=file)
-
-            if not await self.file_exists_in_s3(f"{model_name}/tokenizer.json"):
-                with open(tokenizer_file, "rb") as file:
-                    self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/tokenizer.json", Body=file)
-
+            files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_name)
+            model_files = [obj['Key'] for obj in files.get('Contents', []) if model_name in obj['Key']]
+            return model_files
         except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Error al reanudar la descarga del modelo: {e}")
+            raise HTTPException(status_code=500, detail=f"Error al obtener archivos del modelo {model_name} desde S3: {str(e)}")
 
 @app.post("/generate")
 async def generate(request: GenerateRequest):
@@ -187,7 +142,7 @@ async def generate(request: GenerateRequest):
         
         if task_type == "text-to-text":
             generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)
-            result = generator(input_text, max_length=MAX_TOKENS, num_return_sequences=1)
+            result = generator(input_text, max_length=1024, num_return_sequences=1)
             return {"result": result[0]["generated_text"]}
 
         elif task_type == "text-to-image":