Update app.py

app.py

@@ -1,193 +1,197 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-import requests
-import boto3
-from dotenv import load_dotenv
 import os
-import uvicorn
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
-import torch
-import safetensors.torch
+import logging
+import time
+from io import BytesIO
+from typing import Union
+
+from fastapi import FastAPI, HTTPException, Response, Request, UploadFile, File
 from fastapi.responses import StreamingResponse
-from tqdm import tqdm
+from pydantic import BaseModel, ValidationError, field_validator
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    pipeline,
+    GenerationConfig,
+    StoppingCriteriaList
+)
+import boto3
+from huggingface_hub import hf_hub_download
+import soundfile as sf
+import numpy as np
+import torch
+import uvicorn
 
-# Cargar las variables de entorno desde el archivo .env
-load_dotenv()
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s")
 
-# Cargar las credenciales de AWS y el token de Hugging Face desde las 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")  # Nombre del bucket de S3
-HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")  # Token de Hugging Face
-
-# Cliente S3 de Amazon
-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
-)
+S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
+HUGGINGFACE_HUB_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
 
-app = FastAPI()
-
-# Pydantic Model para el cuerpo de la solicitud del endpoint /download_model/
-class DownloadModelRequest(BaseModel):
+class GenerateRequest(BaseModel):
     model_name: str
-    pipeline_task: str
-    input_text: str
-
-class S3DirectStream:
-    def __init__(self, bucket_name):
-        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
-        )
+    input_text: str = ""
+    task_type: str
+    temperature: float = 1.0
+    max_new_tokens: int = 200
+    stream: bool = False
+    top_p: float = 1.0
+    top_k: int = 50
+    repetition_penalty: float = 1.0
+    num_return_sequences: int = 1
+    do_sample: bool = True
+    chunk_delay: float = 0.0
+    stop_sequences: list[str] = []
+
+    model_config = {"protected_namespaces": ()}
+
+    @field_validator("model_name")
+    def model_name_cannot_be_empty(cls, v):
+        if not v:
+            raise ValueError("model_name cannot be empty.")
+        return v
+
+    @field_validator("task_type")
+    def task_type_must_be_valid(cls, v):
+        valid_types = ["text-to-text", "text-to-image", "text-to-speech", "text-to-video"]
+        if v not in valid_types:
+            raise ValueError(f"task_type must be one of: {valid_types}")
+        return v
+
+class S3ModelLoader:
+    def __init__(self, bucket_name, s3_client):
         self.bucket_name = bucket_name
+        self.s3_client = s3_client
 
-    def stream_from_s3(self, key):
-        try:
-            print(f"Descargando archivo {key} desde S3...")
-            response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
-            return response['Body']  # Devolver el cuerpo directamente para el StreamingResponse
-        except self.s3_client.exceptions.NoSuchKey:
-            raise HTTPException(status_code=404, detail=f"El archivo {key} no existe en el bucket S3.")
+    def _get_s3_uri(self, model_name):
+        return f"s3://{self.bucket_name}/{model_name.replace('/', '-')}"
 
-    def file_exists_in_s3(self, key):
+    async def load_model_and_tokenizer(self, model_name):
+        s3_uri = self._get_s3_uri(model_name)
         try:
-            self.s3_client.head_object(Bucket=self.bucket_name, Key=key)
-            return True
-        except self.s3_client.exceptions.ClientError:
-            return False
-
-    def load_model_from_stream(self, model_prefix):
-        try:
-            print(f"Cargando el modelo {model_prefix} desde S3...")
-            if self.file_exists_in_s3(f"{model_prefix}/config.json") and \
-               (self.file_exists_in_s3(f"{model_prefix}/pytorch_model.bin") or self.file_exists_in_s3(f"{model_prefix}/model.safetensors")):
-                print(f"Modelo {model_prefix} ya existe en S3. No es necesario descargarlo.")
-                return self.load_model_from_existing_s3(model_prefix)
-            
-            print(f"Modelo {model_prefix} no encontrado. Procediendo a descargar...")
-            self.download_and_upload_to_s3(model_prefix)
-            return self.load_model_from_stream(model_prefix)
-        except HTTPException as e:
-            print(f"Error al cargar el modelo: {e}")
-            return None
-
-    def load_model_from_existing_s3(self, model_prefix):
-        # Cargar el modelo y los archivos necesarios desde S3
-        print(f"Cargando los archivos {model_prefix} desde S3...")
-        config_stream = self.stream_from_s3(f"{model_prefix}/config.json")
-        config_data = config_stream.read().decode("utf-8")
-        
-        print(f"Cargando el modelo de lenguaje {model_prefix}...")
-
-        # Verificar si el archivo es un safetensor o un archivo binario
-        if self.file_exists_in_s3(f"{model_prefix}/model.safetensors"):
-            # Usar safetensors si el archivo es de tipo safetensors
-            model_stream = self.stream_from_s3(f"{model_prefix}/model.safetensors")
-            model = AutoModelForCausalLM.from_config(config_data)
-            model.load_state_dict(safetensors.torch.load_stream(model_stream))  # Cargar el modelo utilizando safetensors
-        else:
-            # Cargar el modelo utilizando pytorch si el archivo es .bin
-            model_stream = self.stream_from_s3(f"{model_prefix}/pytorch_model.bin")
-            model = AutoModelForCausalLM.from_config(config_data)
-            model.load_state_dict(torch.load(model_stream, map_location="cpu"))
+            logging.info(f"Trying to load {model_name} from S3...")
+            config = AutoConfig.from_pretrained(s3_uri)
+            model = AutoModelForCausalLM.from_pretrained(s3_uri, config=config)
+            tokenizer = AutoTokenizer.from_pretrained(s3_uri, config=config)
+
+            if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
+                tokenizer.pad_token_id = config.pad_token_id or tokenizer.eos_token_id
+
+            logging.info(f"Loaded {model_name} from S3 successfully.")
+            return model, tokenizer
+        except EnvironmentError:
+            logging.info(f"Model {model_name} not found in S3. Downloading...")
+            try:
+                config = AutoConfig.from_pretrained(model_name)
+                tokenizer = AutoTokenizer.from_pretrained(model_name, config=config)
+                model = AutoModelForCausalLM.from_pretrained(model_name, config=config, token=HUGGINGFACE_HUB_TOKEN)
+
+                if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
+                    tokenizer.pad_token_id = config.pad_token_id or tokenizer.eos_token_id
+
+                logging.info(f"Downloaded {model_name} successfully.")
+                logging.info(f"Saving {model_name} to S3...")
+                model.save_pretrained(s3_uri)
+                tokenizer.save_pretrained(s3_uri)
+                logging.info(f"Saved {model_name} to S3 successfully.")
+                return model, tokenizer
+            except Exception as e:
+                logging.exception(f"Error downloading/uploading model: {e}")
+                raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
 
-        return model
-
-    def load_tokenizer_from_stream(self, model_prefix):
-        try:
-            if self.file_exists_in_s3(f"{model_prefix}/tokenizer.json"):
-                print(f"Tokenizer para {model_prefix} ya existe en S3. No es necesario descargarlo.")
-                return self.load_tokenizer_from_existing_s3(model_prefix)
-            
-            print(f"Tokenizer para {model_prefix} no encontrado. Procediendo a descargar...")
-            self.download_and_upload_to_s3(model_prefix)
-            return self.load_tokenizer_from_stream(model_prefix)
-        except HTTPException as e:
-            print(f"Error al cargar el tokenizer: {e}")
-            return None
-
-    def load_tokenizer_from_existing_s3(self, model_prefix):
-        print(f"Cargando el tokenizer para {model_prefix} desde S3...")
-        tokenizer_stream = self.stream_from_s3(f"{model_prefix}/tokenizer.json")
-        tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
-        return tokenizer
-
-    def download_and_upload_to_s3(self, model_prefix):
-        # URLs de los archivos de Hugging Face
-        model_url = f"https://huggingface.co/{model_prefix}/resolve/main/pytorch_model.bin"
-        safetensors_url = f"https://huggingface.co/{model_prefix}/resolve/main/model.safetensors"
-        tokenizer_url = f"https://huggingface.co/{model_prefix}/resolve/main/tokenizer.json"
-        config_url = f"https://huggingface.co/{model_prefix}/resolve/main/config.json"
-
-        print(f"Descargando y subiendo archivos para el modelo {model_prefix} a S3...")
-        self.download_and_upload_to_s3_url(model_url, f"{model_prefix}/pytorch_model.bin")
-        self.download_and_upload_to_s3_url(safetensors_url, f"{model_prefix}/model.safetensors")
-        self.download_and_upload_to_s3_url(tokenizer_url, f"{model_prefix}/tokenizer.json")
-        self.download_and_upload_to_s3_url(config_url, f"{model_prefix}/config.json")
-
-    def download_and_upload_to_s3_url(self, url: str, s3_key: str):
-        print(f"Descargando archivo desde {url}...")
-        response = requests.get(url)
-        if response.status_code == 200:
-            # Subir archivo a S3
-            print(f"Subiendo archivo a S3 con key {s3_key}...")
-            self.s3_client.put_object(Bucket=self.bucket_name, Key=s3_key, Body=response.content)
-        else:
-            raise HTTPException(status_code=500, detail=f"Error al descargar el archivo desde {url}")
+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)
 
-@app.post("/predict/")
-async def predict(model_request: DownloadModelRequest):
+@app.post("/generate")
+async def generate(request: Request, body: GenerateRequest):
     try:
-        print(f"Recibiendo solicitud para predecir con el modelo {model_request.model_name}...")
-        # Cargar el modelo y tokenizer desde S3
-        streamer = S3DirectStream(S3_BUCKET_NAME)
-        model = streamer.load_model_from_stream(model_request.model_name)
-        tokenizer = streamer.load_tokenizer_from_stream(model_request.model_name)
-
-        # Obtener el pipeline adecuado según la solicitud
-        task = model_request.pipeline_task
-        if task not in ["text-generation", "sentiment-analysis", "translation", "fill-mask", "question-answering", "text-to-speech", "text-to-image", "text-to-audio", "text-to-video"]:
-            raise HTTPException(status_code=400, detail="Pipeline task no soportado")
-
-        # Crear el pipeline dinámicamente basado en el tipo de tarea
-        nlp_pipeline = pipeline(task, model=model, tokenizer=tokenizer)
-
-        # Ejecutar el pipeline con el input_text
-        input_text = model_request.input_text
-        outputs = nlp_pipeline(input_text)
-
-        # Procesar los diferentes tipos de respuestas según el pipeline
-        if task in ["text-generation", "translation", "fill-mask", "sentiment-analysis", "question-answering"]:
-            return {"response": outputs}
-
-        elif task == "text-to-image":
-            # Asumir que outputs es la imagen generada
-            s3_key = f"{model_request.model_name}/generated_image.png"  # Definir el key del archivo de imagen
-            return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="image/png")
-
-        elif task == "text-to-audio":
-            # Asumir que outputs es el audio generado
-            s3_key = f"{model_request.model_name}/generated_audio.wav"  # Definir el key del archivo de audio
-            return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="audio/wav")
-
-        elif task == "text-to-video":
-            # Asumir que outputs es el video generado
-            s3_key = f"{model_request.model_name}/generated_video.mp4"  # Definir el key del archivo de video
-            return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="video/mp4")
+        validated_body = GenerateRequest(**body.model_dump())
+        model, tokenizer = await model_loader.load_model_and_tokenizer(validated_body.model_name)
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        model.to(device)
+
+        if validated_body.task_type == "text-to-text":
+            generation_config = GenerationConfig(
+                temperature=validated_body.temperature,
+                max_new_tokens=validated_body.max_new_tokens,
+                top_p=validated_body.top_p,
+                top_k=validated_body.top_k,
+                repetition_penalty=validated_body.repetition_penalty,
+                do_sample=validated_body.do_sample,
+                num_return_sequences=validated_body.num_return_sequences
+            )
+
+            async def stream_text():
+                input_text = validated_body.input_text
+                generated_text = ""
+                max_length = model.config.max_position_embeddings
+
+                while True:
+                    encoded_input = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=max_length).to(device)
+                    input_length = encoded_input["input_ids"].shape[1]
+                    remaining_tokens = max_length - input_length
+
+                    if remaining_tokens <= 0:
+                        break
+
+                    generation_config.max_new_tokens = min(remaining_tokens, validated_body.max_new_tokens)
+
+                    stopping_criteria = StoppingCriteriaList(
+                        [lambda _, outputs: tokenizer.decode(outputs[0][-1], skip_special_tokens=True) in validated_body.stop_sequences] if validated_body.stop_sequences else []
+                    )
+
+                    output = model.generate(**encoded_input, generation_config=generation_config, stopping_criteria=stopping_criteria)
+                    chunk = tokenizer.decode(output[0], skip_special_tokens=True)
+                    generated_text += chunk
+                    yield chunk
+                    time.sleep(validated_body.chunk_delay)
+                    input_text = generated_text
+
+            if validated_body.stream:
+                return StreamingResponse(stream_text(), media_type="text/plain")
+            else:
+                generated_text = ""
+                async for chunk in stream_text():
+                    generated_text += chunk
+                return {"result": generated_text}
+
+        elif validated_body.task_type == "text-to-image":
+            generator = pipeline("text-to-image", model=model, tokenizer=tokenizer, device=device)
+            image = generator(validated_body.input_text)[0]
+            image_bytes = image.tobytes()
+            return Response(content=image_bytes, media_type="image/png")
+
+        elif validated_body.task_type == "text-to-speech":
+            generator = pipeline("text-to-speech", model=model, tokenizer=tokenizer, device=device)
+            audio = generator(validated_body.input_text)
+            audio_bytesio = BytesIO()
+            sf.write(audio_bytesio, audio["sampling_rate"], np.int16(audio["audio"]))
+            audio_bytes = audio_bytesio.getvalue()
+            return Response(content=audio_bytes, media_type="audio/wav")
+
+        elif validated_body.task_type == "text-to-video":
+            try:
+                generator = pipeline("text-to-video", model=model, tokenizer=tokenizer, device=device)
+                video = generator(validated_body.input_text)
+                return Response(content=video, media_type="video/mp4")
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=f"Error in text-to-video generation: {e}")
 
         else:
-            raise HTTPException(status_code=400, detail="Tipo de tarea desconocido")
+            raise HTTPException(status_code=400, detail="Unsupported task type")
 
+    except HTTPException as e:
+        raise e
+    except ValidationError as e:
+        raise HTTPException(status_code=422, detail=e.errors())
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error al procesar la solicitud: {str(e)}")
+        logging.exception(f"An unexpected error occurred: {e}")
+        raise HTTPException(status_code=500, detail="An unexpected error occurred.")
 
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)