import os
import boto3
import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import safetensors.torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from fastapi.responses import StreamingResponse
import io
import requests
import uvicorn
from dotenv import load_dotenv
import re
from tqdm import tqdm
from huggingface_hub import hf_hub_download
# 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 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.")
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_stream(self, model_name):
try:
print(f"[INFO] Cargando el modelo {model_name} desde 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.")
# Leer y cargar todos los archivos del modelo
model_streams = []
for model_file in 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 = config_stream.read().decode("utf-8")
# Cargar el modelo dependiendo del tipo de archivo (torch o safetensors)
if model_files[0].endswith("model.safetensors"):
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_stream(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(files.get('Contents', []), desc="Verificando archivos", unit="archivo"):
key = obj['Key']
if re.match(rf"{model_name}/model(-\d+-of-\d+)?", key) or key.endswith("model.safetensors"):
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 los archivos necesarios de Hugging Face usando huggingface_hub
files_to_download = [
"pytorch_model.bin",
"config.json",
"tokenizer.json",
]
for file in files_to_download:
# Descargar cada archivo desde Hugging Face y subirlo a S3
file_path = hf_hub_download(repo_id=model_name, filename=file, use_auth_token=HUGGINGFACE_TOKEN)
s3_key = f"{model_name}/{file}"
self.upload_file_to_s3(file_path, 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)
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_stream(model_request.model_name)
tokenizer = streamer.load_tokenizer_from_stream(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-image", "text-to-audio", "text-to-video"]:
raise HTTPException(status_code=400, detail="Pipeline task no soportado")
nlp_pipeline = pipeline(task, model=model, tokenizer=tokenizer, max_length=2046)
input_text = model_request.input_text
print(f"[INFO] Ejecutando tarea {task} con el texto de entrada...")
outputs = nlp_pipeline(input_text)
# Eliminación de archivo local después de subir a S3
if task == "text-to-speech":
s3_key = f"{model_request.model_name}/generated_audio.wav"
return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="audio/wav")
elif task == "text-to-image":
s3_key = f"{model_request.model_name}/generated_image.png"
return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="image/png")
elif task == "text-to-video":
s3_key = f"{model_request.model_name}/generated_video.mp4"
return StreamingResponse(streamer.stream_from_s3(s3_key), media_type="video/mp4")
return {"output": outputs}
except Exception as e:
print(f"[ERROR] Error al procesar la solicitud: {str(e)}")
raise HTTPException(status_code=500, detail="Error al realizar la predicción.")
if __name__ == "__main__":
print("Iniciando servidor de predicciones en localhost:8000")
uvicorn.run(app, host="0.0.0.0", port=7860)