import os import json import logging import boto3 from fastapi import FastAPI, HTTPException, Query from fastapi.responses import JSONResponse from transformers import AutoModelForCausalLM, AutoTokenizer from huggingface_hub import hf_hub_download import asyncio logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) console_handler = logging.StreamHandler() formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') console_handler.setFormatter(formatter) logger.addHandler(console_handler) 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") MAX_TOKENS = 1024 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() 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 async def stream_from_s3(self, key): loop = asyncio.get_event_loop() return await loop.run_in_executor(None, self._stream_from_s3, key) def _stream_from_s3(self, key): try: response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key) return response['Body'] 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: raise HTTPException(status_code=500, detail=f"Error al descargar {key} desde S3: {str(e)}") async def get_model_file_parts(self, model_name): loop = asyncio.get_event_loop() return await loop.run_in_executor(None, self._get_model_file_parts, model_name) def _get_model_file_parts(self, model_name): try: model_name = model_name.replace("/", "-").lower() files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_name) model_files = [obj['Key'] for obj in files.get('Contents', []) if model_name in obj['Key']] return model_files except Exception as e: raise HTTPException(status_code=500, detail=f"Error al obtener archivos del modelo {model_name} desde S3: {e}") async def load_model_from_s3(self, model_name): try: model_name = model_name.replace("/", "-").lower() model_files = await self.get_model_file_parts(model_name) if not model_files: await self.download_and_upload_to_s3(model_name) config_stream = await self.stream_from_s3(f"{model_name}/config.json") config_data = config_stream.read() if not config_data: raise HTTPException(status_code=500, detail=f"El archivo de configuración {model_name}/config.json está vacío o no se pudo leer.") config_text = config_data.decode("utf-8") config_json = json.loads(config_text) model = AutoModelForCausalLM.from_pretrained(f"s3://{self.bucket_name}/{model_name}", config=config_json, from_tf=False) return model except HTTPException as e: raise e except Exception as e: raise HTTPException(status_code=500, detail=f"Error al cargar el modelo desde S3: {e}") async def load_tokenizer_from_s3(self, model_name): try: model_name = model_name.replace("/", "-").lower() tokenizer_stream = await self.stream_from_s3(f"{model_name}/tokenizer.json") tokenizer_data = tokenizer_stream.read().decode("utf-8") tokenizer = AutoTokenizer.from_pretrained(f"s3://{self.bucket_name}/{model_name}") return tokenizer except Exception as e: raise HTTPException(status_code=500, detail=f"Error al cargar el tokenizer desde S3: {e}") async def create_s3_folders(self, s3_key): try: folder_keys = s3_key.split('-') for i in range(1, len(folder_keys)): folder_key = '-'.join(folder_keys[:i]) + '/' if not await self.file_exists_in_s3(folder_key): logger.info(f"Creando carpeta en S3: {folder_key}") self.s3_client.put_object(Bucket=self.bucket_name, Key=folder_key, Body='') except Exception as e: raise HTTPException(status_code=500, detail=f"Error al crear carpetas en S3: {e}") async def file_exists_in_s3(self, s3_key): try: self.s3_client.head_object(Bucket=self.bucket_name, Key=s3_key) return True except self.s3_client.exceptions.ClientError: return False async def download_and_upload_to_s3(self, model_name, force_download=False): try: if force_download: logger.info(f"Forzando la descarga del modelo {model_name} y la carga a S3.") model_name = model_name.replace("/", "-").lower() if not await self.file_exists_in_s3(f"{model_name}/config.json") or not await self.file_exists_in_s3(f"{model_name}/tokenizer.json"): config_file = hf_hub_download(repo_id=model_name, filename="config.json", token=HUGGINGFACE_HUB_TOKEN, force_download=force_download) tokenizer_file = hf_hub_download(repo_id=model_name, filename="tokenizer.json", token=HUGGINGFACE_HUB_TOKEN, force_download=force_download) await self.create_s3_folders(f"{model_name}/") if not await self.file_exists_in_s3(f"{model_name}/config.json"): with open(config_file, "rb") as file: self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/config.json", Body=file) if not await self.file_exists_in_s3(f"{model_name}/tokenizer.json"): with open(tokenizer_file, "rb") as file: self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/tokenizer.json", Body=file) else: logger.info(f"Los archivos del modelo {model_name} ya existen en S3. No es necesario descargarlos de nuevo.") except Exception as e: raise HTTPException(status_code=500, detail=f"Error al descargar o cargar archivos desde Hugging Face a S3: {e}") async def resume_download(self, model_name): try: logger.info(f"Reanudando la descarga del modelo {model_name} desde Hugging Face.") config_file = hf_hub_download(repo_id=model_name, filename="config.json", token=HUGGINGFACE_HUB_TOKEN, resume_download=True) tokenizer_file = hf_hub_download(repo_id=model_name, filename="tokenizer.json", token=HUGGINGFACE_HUB_TOKEN, resume_download=True) if not await self.file_exists_in_s3(f"{model_name}/config.json"): with open(config_file, "rb") as file: self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/config.json", Body=file) if not await self.file_exists_in_s3(f"{model_name}/tokenizer.json"): with open(tokenizer_file, "rb") as file: self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/tokenizer.json", Body=file) except Exception as e: raise HTTPException(status_code=500, detail=f"Error al reanudar la descarga o cargar archivos desde Hugging Face a S3: {e}") def split_text_by_tokens(text, tokenizer, max_tokens=MAX_TOKENS): tokens = tokenizer.encode(text) chunks = [] for i in range(0, len(tokens), max_tokens): chunk = tokens[i:i+max_tokens] chunks.append(tokenizer.decode(chunk)) return chunks def continue_generation(input_text, model, tokenizer, max_tokens=MAX_TOKENS): generated_text = "" while len(input_text) > 0: chunks = split_text_by_tokens(input_text, tokenizer, max_tokens) for chunk in chunks: generated_text += model.generate(chunk) return generated_text @app.post("/generate") async def generate_text(model_name: str = Query(...), input_text: str = Query(...)): try: model_loader = S3DirectStream(S3_BUCKET_NAME) model = await model_loader.load_model_from_s3(model_name) tokenizer = await model_loader.load_tokenizer_from_s3(model_name) chunks = split_text_by_tokens(input_text, tokenizer, max_tokens=MAX_TOKENS) generated_text = continue_generation(input_text, model, tokenizer) return {"generated_text": generated_text} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)