app.py

import os
import json
import boto3
import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from huggingface_hub import hf_hub_download
from io import BytesIO
import torch
import safetensors
from dotenv import load_dotenv
import tqdm
import re

# Cargar las variables de entorno desde el archivo .env
load_dotenv()

# Configuración de AWS y Hugging Face
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 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
)

app = FastAPI()

# Pydantic Model para el cuerpo de la solicitud del endpoint /predict/
class DownloadModelRequest(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
        )
        self.bucket_name = bucket_name

    def stream_from_s3(self, key):
        try:
            print(f"[INFO] Descargando archivo {key} desde S3...")
            response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
            return BytesIO(response['Body'].read())  # Devolver el cuerpo como BytesIO
        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:
            print(f"[ERROR] Error al descargar {key}: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Error al descargar archivo {key} desde S3.")

    def file_exists_in_s3(self, key):
        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_s3(self, model_name):
        try:
            print(f"[INFO] Cargando el modelo {model_name} desde S3...")

            # Verificar si el modelo existe en S3
            model_prefix = model_name.lower()
            model_files = self.get_model_file_parts(model_prefix)

            if not model_files:
                print(f"[INFO] El modelo {model_name} no está en S3. Procediendo a descargar desde Hugging Face...")
                self.download_and_upload_from_huggingface(model_name)
                model_files = self.get_model_file_parts(model_prefix)

            if not model_files:
                raise HTTPException(status_code=404, detail=f"Archivos del modelo {model_name} no encontrados en S3.")

            # Cargar todos los archivos del modelo desde S3
            model_streams = []
            for model_file in tqdm.tqdm(model_files, desc="Cargando archivos del modelo", unit="archivo"):
                model_streams.append(self.stream_from_s3(model_file))

            config_stream = self.stream_from_s3(f"{model_prefix}/config.json")
            config_data = json.loads(config_stream.read().decode("utf-8"))

            # Cargar el modelo dependiendo del tipo de archivo (torch o safetensors)
            if any(file.endswith("model.safetensors") for file in model_files):
                print("[INFO] Cargando el modelo como safetensor...")
                model = AutoModelForCausalLM.from_config(config_data)
                model.load_state_dict(safetensors.torch.load_stream(model_streams[0]))
            else:
                print("[INFO] Cargando el modelo como archivo binario de PyTorch...")
                model = AutoModelForCausalLM.from_config(config_data)
                model.load_state_dict(torch.load(model_streams[0], map_location="cpu"))

            print("[INFO] Modelo cargado con éxito.")
            return model

        except Exception as e:
            print(f"[ERROR] Error al cargar el modelo desde S3: {e}")
            raise HTTPException(status_code=500, detail="Error al cargar el modelo desde S3.")

    def load_tokenizer_from_s3(self, model_name):
        try:
            print(f"[INFO] Cargando el tokenizer {model_name} desde S3...")
            tokenizer_stream = self.stream_from_s3(f"{model_name}/tokenizer.json")
            tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
            return tokenizer
        except Exception as e:
            print(f"[ERROR] Error al cargar el tokenizer desde S3: {e}")
            raise HTTPException(status_code=500, detail="Error al cargar el tokenizer desde S3.")

    def get_model_file_parts(self, model_name):
        print(f"[INFO] Listando archivos del modelo en S3 con prefijo {model_name}...")
        files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_name)
        model_files = []
        for obj in tqdm.tqdm(files.get('Contents', []), desc="Verificando archivos", unit="archivo"):
            key = obj['Key']
            # Verificar si es un archivo relevante del modelo
            if re.match(rf"{model_name}/.*", key):
                model_files.append(key)
        if not model_files:
            print(f"[WARNING] No se encontraron archivos para el modelo {model_name}.")
        return model_files

    def download_and_upload_from_huggingface(self, model_name):
        try:
            print(f"[INFO] Descargando {model_name} desde Hugging Face...")
            # Descargar todos los archivos del modelo
            files_to_download = hf_hub_download(repo_id=model_name, use_auth_token=HUGGINGFACE_TOKEN)

            # Subir a S3 los archivos descargados
            for file in files_to_download:
                file_name = os.path.basename(file)
                s3_key = f"{model_name}/{file_name}"
                if not self.file_exists_in_s3(s3_key):
                    self.upload_file_to_s3(file, s3_key)

        except Exception as e:
            print(f"[ERROR] Error al descargar y subir modelo desde Hugging Face: {e}")
            raise HTTPException(status_code=500, detail="Error al descargar y subir modelo desde Hugging Face.")

    def upload_file_to_s3(self, file_path, s3_key):
        try:
            print(f"[INFO] Subiendo archivo {file_path} a S3 con key {s3_key}...")
            with open(file_path, 'rb') as data:
                self.s3_client.put_object(Bucket=self.bucket_name, Key=s3_key, Body=data)
            # Eliminar archivo local después de ser subido
            os.remove(file_path)
        except Exception as e:
            print(f"[ERROR] Error al subir archivo a S3: {e}")
            raise HTTPException(status_code=500, detail="Error al subir archivo a S3.")


@app.post("/predict/")
async def predict(model_request: DownloadModelRequest):
    try:
        print(f"[INFO] Recibiendo solicitud para predecir con el modelo {model_request.model_name}...")
        streamer = S3DirectStream(S3_BUCKET_NAME)
        model = streamer.load_model_from_s3(model_request.model_name)
        tokenizer = streamer.load_tokenizer_from_s3(model_request.model_name)

        task = model_request.pipeline_task
        if task not in ["text-generation", "sentiment-analysis", "translation", "fill-mask", "question-answering", 
                        "text-to-speech", "text-to-video", "text-to-image"]:
            raise HTTPException(status_code=400, detail="Pipeline task no soportado")

        # Configurar el pipeline de transformers según la tarea
        nlp_pipeline = None
        if task == "text-generation":
            nlp_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
        elif task == "sentiment-analysis":
            nlp_pipeline = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer)
        elif task == "translation":
            nlp_pipeline = pipeline("translation", model=model, tokenizer=tokenizer)
        elif task == "fill-mask":
            nlp_pipeline = pipeline("fill-mask", model=model, tokenizer=tokenizer)
        elif task == "question-answering":
            nlp_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)
        elif task == "text-to-speech":
            nlp_pipeline = pipeline("text-to-speech", model=model, tokenizer=tokenizer)
        elif task == "text-to-video":
            nlp_pipeline = pipeline("text-to-video", model=model, tokenizer=tokenizer)
        elif task == "text-to-image":
            nlp_pipeline = pipeline("text-to-image", model=model, tokenizer=tokenizer)

        result = nlp_pipeline(model_request.input_text)
        return {"result": result}

    except Exception as e:
        print(f"[ERROR] Error en el proceso de predicción: {str(e)}")
        raise HTTPException(status_code=500, detail="Error en el proceso de predicción")


# Ejecutar la app con Uvicorn
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)