Update app.py

app.py

@@ -3,32 +3,31 @@ import logging
 import requests
 import threading
 from io import BytesIO
-from fastapi import FastAPI, HTTPException, Response, Request
+from fastapi import FastAPI, HTTPException, Response
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from transformers import (
     AutoConfig,
     AutoModelForCausalLM,
     AutoTokenizer,
-    pipeline,
     GenerationConfig
 )
 import boto3
-from huggingface_hub import hf_hub_download
-import soundfile as sf
-import numpy as np
 import torch
 import uvicorn
 from tqdm import tqdm
 
+# Configuración de logging
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
+# Variables de entorno
 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")
 
+# Clase para la petición de generación
 class GenerateRequest(BaseModel):
     model_name: str
     input_text: str
@@ -42,6 +41,10 @@ class GenerateRequest(BaseModel):
     num_return_sequences: int = 1
     do_sample: bool = True
 
+    class Config:
+        protected_namespaces = ()
+
+# Clase para cargar modelos desde S3
 class S3ModelLoader:
     def __init__(self, bucket_name, s3_client):
         self.bucket_name = bucket_name
@@ -71,8 +74,8 @@ class S3ModelLoader:
     async def download_and_save_model_from_huggingface(self, model_name):
         try:
             with tqdm(unit="B", unit_scale=True, desc=f"Downloading {model_name}") as t:
-                model = AutoModelForCausalLM.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN, _tqdm=t)
-                tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN)
+                model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=HUGGINGFACE_HUB_TOKEN, _tqdm=t)
+                tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=HUGGINGFACE_HUB_TOKEN)
             self.upload_model_to_s3(model_name, model, tokenizer)
             return model, tokenizer
         except Exception as e:
@@ -86,13 +89,12 @@ class S3ModelLoader:
         except Exception as e:
             raise HTTPException(status_code=500, detail=f"Error saving model to S3: {e}")
 
+# Crear la instancia de FastAPI
 app = FastAPI()
 
-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)
-model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)
-
+# Función de generación asincrónica
 @app.post("/generate")
-async def generate(request: Request, body: GenerateRequest):
+async def generate(body: GenerateRequest):
     try:
         model, tokenizer = await model_loader.load_model_and_tokenizer(body.model_name)
         device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -169,7 +171,8 @@ async def generate(request: Request, body: GenerateRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-def download_all_models_in_background():
+# Descargar todos los modelos en segundo plano
+async def download_all_models_in_background():
     models_url = "https://huggingface.co/api/models"
     try:
         response = requests.get(models_url)
@@ -179,16 +182,14 @@ def download_all_models_in_background():
         models = response.json()
         for model in models:
             model_name = model["id"]
-            model_loader.download_and_save_model_from_huggingface(model_name)
+            await model_loader.download_and_save_model_from_huggingface(model_name)
     except Exception as e:
         raise HTTPException(status_code=500, detail="Error al descargar modelos en segundo plano.")
 
+# Función que corre en segundo plano para descargar modelos
 def run_in_background():
     threading.Thread(target=download_all_models_in_background, daemon=True).start()
 
-@app.on_event("startup")
-async def startup_event():
-    run_in_background()
-
+# Si este archivo se ejecuta directamente, inicia el servidor
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)