app.py.bkk

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)