Update app.py

app.py

@@ -5,27 +5,31 @@ from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from transformers import (
     AutoConfig,
-    pipeline,
     AutoModelForCausalLM,
     AutoTokenizer,
     GenerationConfig,
-    StoppingCriteriaList
+    StoppingCriteriaList,
+    pipeline
 )
 import asyncio
 from io import BytesIO
+from botocore.exceptions import NoCredentialsError
+import boto3
 
-# Diccionario global para almacenar los tokens
+# Diccionario global para almacenar los tokens y configuraciones de los modelos
 token_dict = {}
 
-# Setup para acceder a modelos en Hugging Face o S3
+# Configuración para acceso a modelos en Hugging Face o S3
 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")
 
+# Inicialización de la aplicación FastAPI
 app = FastAPI()
 
+# Modelo de la solicitud para la API
 class GenerateRequest(BaseModel):
     model_name: str
     input_text: str
@@ -42,14 +46,19 @@ class GenerateRequest(BaseModel):
     stop_sequences: list[str] = []
 
 class S3ModelLoader:
-    def __init__(self, bucket_name, s3_client):
+    def __init__(self, bucket_name, aws_access_key_id=None, aws_secret_access_key=None, aws_region=None):
         self.bucket_name = bucket_name
-        self.s3_client = s3_client
+        self.s3_client = boto3.client(
+            's3',
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+            region_name=aws_region
+        )
 
     def _get_s3_uri(self, model_name):
         return f"s3://{self.bucket_name}/{model_name.replace('/', '-')}"
     
-    async def load_model_and_tokenizer(self, model_name):
+    def load_model_and_tokenizer(self, model_name):
         if model_name in token_dict:
             return token_dict[model_name]
         
@@ -69,55 +78,14 @@ class S3ModelLoader:
             }
 
             return token_dict[model_name]
+        except NoCredentialsError:
+            raise HTTPException(status_code=500, detail="AWS credentials not found.")
         except Exception as e:
             raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
 
-model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)
-
-@app.post("/generate")
-async def generate(request: GenerateRequest):
-    try:
-        model_name = request.model_name
-        input_text = request.input_text
-        temperature = request.temperature
-        max_new_tokens = request.max_new_tokens
-        stream = request.stream
-        top_p = request.top_p
-        top_k = request.top_k
-        repetition_penalty = request.repetition_penalty
-        num_return_sequences = request.num_return_sequences
-        do_sample = request.do_sample
-        chunk_delay = request.chunk_delay
-        stop_sequences = request.stop_sequences
-
-        # Cargar modelo y tokenizer desde el S3
-        model_data = await model_loader.load_model_and_tokenizer(model_name)
-        model = model_data["model"]
-        tokenizer = model_data["tokenizer"]
-        pad_token_id = model_data["pad_token_id"]
-        eos_token_id = model_data["eos_token_id"]
-
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        model.to(device)
-
-        generation_config = GenerationConfig(
-            temperature=temperature,
-            max_new_tokens=max_new_tokens,
-            top_p=top_p,
-            top_k=top_k,
-            repetition_penalty=repetition_penalty,
-            do_sample=do_sample,
-            num_return_sequences=num_return_sequences,
-        )
-
-        return StreamingResponse(
-            stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay),
-            media_type="text/plain"
-        )
-
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+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
 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]
@@ -159,20 +127,52 @@ async def stream_text(model, tokenizer, input_text, generation_config, stop_sequ
             yield output_text
             return
 
-        outputs = model.generate(
-            **encoded_input,
-            do_sample=generation_config.do_sample,
-            max_new_tokens=generation_config.max_new_tokens,
-            temperature=generation_config.temperature,
-            top_p=generation_config.top_p,
-            top_k=generation_config.top_k,
-            repetition_penalty=generation_config.repetition_penalty,
-            num_return_sequences=generation_config.num_return_sequences,
-            stopping_criteria=stopping_criteria,
-            output_scores=True,
-            return_dict_in_generate=True
+# Endpoint para la generación de texto
+@app.post("/generate")
+async def generate(request: GenerateRequest):
+    try:
+        model_name = request.model_name
+        input_text = request.input_text
+        temperature = request.temperature
+        max_new_tokens = request.max_new_tokens
+        stream = request.stream
+        top_p = request.top_p
+        top_k = request.top_k
+        repetition_penalty = request.repetition_penalty
+        num_return_sequences = request.num_return_sequences
+        do_sample = request.do_sample
+        chunk_delay = request.chunk_delay
+        stop_sequences = request.stop_sequences
+
+        # Cargar el modelo y el tokenizer desde el S3
+        model_data = model_loader.load_model_and_tokenizer(model_name)
+        model = model_data["model"]
+        tokenizer = model_data["tokenizer"]
+        pad_token_id = model_data["pad_token_id"]
+        eos_token_id = model_data["eos_token_id"]
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        model.to(device)
+
+        generation_config = GenerationConfig(
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            top_p=top_p,
+            top_k=top_k,
+            repetition_penalty=repetition_penalty,
+            do_sample=do_sample,
+            num_return_sequences=num_return_sequences,
+        )
+
+        return StreamingResponse(
+            stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay),
+            media_type="text/plain"
         )
 
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+
+# Endpoint para la generación de imágenes
 @app.post("/generate-image")
 async def generate_image(request: GenerateRequest):
     try:
@@ -191,6 +191,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
 @app.post("/generate-text-to-speech")
 async def generate_text_to_speech(request: GenerateRequest):
     try:
@@ -209,6 +210,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
 @app.post("/generate-video")
 async def generate_video(request: GenerateRequest):
     try:
@@ -226,6 +228,7 @@ async def generate_video(request: GenerateRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
+# Configuración para ejecutar el servidor
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)