Update app.py

app.py

@@ -10,26 +10,26 @@ from transformers import (
     StoppingCriteriaList,
     pipeline
 )
-import asyncio
 from io import BytesIO
-from botocore.exceptions import NoCredentialsError
+import asyncio
 import boto3
+from botocore.exceptions import NoCredentialsError
 from huggingface_hub import snapshot_download
 
-# Diccionario global para almacenar los tokens y configuraciones de los modelos
-token_dict = {}
-
-# Configuración para acceso a modelos en Hugging Face o S3
+# Configuración global
 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")
 
+# Diccionario global de tokens y configuraciones
+token_dict = {}
+
 # Inicialización de la aplicación FastAPI
 app = FastAPI()
 
-# Modelo de la solicitud para la API
+# Modelo de solicitud
 class GenerateRequest(BaseModel):
     model_name: str
     input_text: str
@@ -45,6 +45,7 @@ class GenerateRequest(BaseModel):
     chunk_delay: float = 0.0
     stop_sequences: list[str] = []
 
+# Clase para cargar y gestionar los modelos desde S3
 class S3ModelLoader:
     def __init__(self, bucket_name, aws_access_key_id=None, aws_secret_access_key=None, aws_region=None):
         self.bucket_name = bucket_name
@@ -64,16 +65,18 @@ class S3ModelLoader:
         
         s3_uri = self._get_s3_uri(model_name)
         try:
-            # Descargamos el modelo y el tokenizer desde Hugging Face directamente a S3
+            # Descargar el modelo desde Hugging Face y guardarlo en S3 si no existe
             model_path = snapshot_download(model_name, token=HUGGINGFACE_HUB_TOKEN)
-            
+
+            # Cargar el modelo y tokenizer
             model = AutoModelForCausalLM.from_pretrained(model_path)
             tokenizer = AutoTokenizer.from_pretrained(model_path)
 
+            # Asignar EOS y PAD token si no están definidos
             if tokenizer.eos_token_id is None:
                 tokenizer.eos_token_id = tokenizer.pad_token_id
 
-            # Guardamos en el diccionario global
+            # Guardar el modelo y el tokenizer en el diccionario
             token_dict[model_name] = {
                 "model": model,
                 "tokenizer": tokenizer,
@@ -81,7 +84,7 @@ class S3ModelLoader:
                 "eos_token_id": tokenizer.eos_token_id
             }
 
-            # Subimos los modelos al S3 si es necesario
+            # Subir los archivos del modelo y tokenizer a S3 si no están allí
             self.s3_client.upload_file(model_path, self.bucket_name, f'{model_name}/model')
             self.s3_client.upload_file(f'{model_path}/tokenizer', self.bucket_name, f'{model_name}/tokenizer')
 
@@ -91,9 +94,10 @@ class S3ModelLoader:
         except Exception as e:
             raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
 
+# Instanciación del cargador de modelos
 model_loader = S3ModelLoader(S3_BUCKET_NAME, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION)
 
-# Función para hacer streaming de texto, generando un token a la vez
+# Función de generación de texto con streaming
 async def stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay, max_length=2048):
     encoded_input = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=max_length).to(device)
     input_length = encoded_input["input_ids"].shape[1]
@@ -129,13 +133,13 @@ async def stream_text(model, tokenizer, input_text, generation_config, stop_sequ
         for token_id in output:
             token = tokenizer.decode(token_id, skip_special_tokens=True)
             yield token
-            await asyncio.sleep(chunk_delay)  # Simula el delay entre tokens
+            await asyncio.sleep(chunk_delay)
 
         if stop_sequences and any(stop in output_text for stop in stop_sequences):
             yield output_text
             return
 
-# Endpoint para la generación de texto
+# Endpoint para generar texto
 @app.post("/generate")
 async def generate(request: GenerateRequest):
     try:
@@ -152,7 +156,7 @@ async def generate(request: GenerateRequest):
         chunk_delay = request.chunk_delay
         stop_sequences = request.stop_sequences
 
-        # Cargar el modelo y el tokenizer desde el S3
+        # Cargar el modelo y tokenizer desde S3 si no existe
         model_data = model_loader.load_model_and_tokenizer(model_name)
         model = model_data["model"]
         tokenizer = model_data["tokenizer"]
@@ -180,7 +184,7 @@ async def generate(request: GenerateRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
-# Endpoint para la generación de imágenes
+# Endpoint para generar imágenes
 @app.post("/generate-image")
 async def generate_image(request: GenerateRequest):
     try:
@@ -199,7 +203,7 @@ async def generate_image(request: GenerateRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
-# Endpoint para la generación de texto a voz
+# Endpoint para generar texto a voz
 @app.post("/generate-text-to-speech")
 async def generate_text_to_speech(request: GenerateRequest):
     try:
@@ -218,7 +222,7 @@ async def generate_text_to_speech(request: GenerateRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
-# Endpoint para la generación de video
+# Endpoint para generar video
 @app.post("/generate-video")
 async def generate_video(request: GenerateRequest):
     try: