Create app.py

app.py

@@ -0,0 +1,218 @@
+import os
+import json
+import logging
+import boto3
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import JSONResponse
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from huggingface_hub import hf_hub_download
+import asyncio
+
+# Configuración del logger
+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()
+
+PIPELINE_MAP = {
+    "text-generation": "text-generation",
+    "sentiment-analysis": "sentiment-analysis",
+    "translation": "translation",
+    "fill-mask": "fill-mask",
+    "question-answering": "question-answering",
+    "text-to-speech": "text-to-speech",
+    "text-to-video": "text-to-video",
+    "text-to-image": "text-to-image"
+}
+
+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_prefix = model_name.lower()
+            files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_prefix)
+            model_files = [obj['Key'] for obj in files.get('Contents', []) if model_prefix 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:
+            profile, model = model_name.split("/", 1) if "/" in model_name else ("", model_name)
+
+            model_prefix = f"{profile}/{model}".lower()
+            model_files = await self.get_model_file_parts(model_prefix)
+
+            if not model_files:
+                await self.download_and_upload_to_s3(model_prefix, model)
+
+            config_stream = await self.stream_from_s3(f"{model_prefix}/config.json")
+            config_data = config_stream.read()
+
+            if not config_data:
+                raise HTTPException(status_code=500, detail=f"El archivo de configuración {model_prefix}/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_prefix}", 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:
+            profile, model = model_name.split("/", 1) if "/" in model_name else ("", model_name)
+
+            tokenizer_stream = await self.stream_from_s3(f"{profile}/{model}/tokenizer.json")
+            tokenizer_data = tokenizer_stream.read().decode("utf-8")
+
+            tokenizer = AutoTokenizer.from_pretrained(f"s3://{self.bucket_name}/{profile}/{model}")
+            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_prefix, model_name):
+        try:
+            config_file = hf_hub_download(repo_id=model_name, filename="config.json", token=HUGGINGFACE_HUB_TOKEN)
+            tokenizer_file = hf_hub_download(repo_id=model_name, filename="tokenizer.json", token=HUGGINGFACE_HUB_TOKEN)
+
+            if not await self.file_exists_in_s3(f"{model_prefix}/config.json"):
+                with open(config_file, "rb") as file:
+                    self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_prefix}/config.json", Body=file)
+
+            if not await self.file_exists_in_s3(f"{model_prefix}/tokenizer.json"):
+                with open(tokenizer_file, "rb") as file:
+                    self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_prefix}/tokenizer.json", Body=file)
+
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"Error al descargar 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:
+        tokens = tokenizer.encode(input_text)
+        input_text = tokenizer.decode(tokens[:max_tokens])
+        output = model.generate(input_ids=tokenizer.encode(input_text, return_tensors="pt").input_ids)
+        generated_text += tokenizer.decode(output[0], skip_special_tokens=True)
+        input_text = input_text[len(input_text):]
+    return generated_text
+
+@app.post("/predict/")
+async def predict(model_request: dict):
+    try:
+        model_name = model_request.get("model_name")
+        task = model_request.get("pipeline_task")
+        input_text = model_request.get("input_text")
+
+        if not model_name or not task or not input_text:
+            raise HTTPException(status_code=400, detail="Faltan parámetros en la solicitud.")
+
+        streamer = S3DirectStream(S3_BUCKET_NAME)
+        
+        await streamer.create_s3_folders(model_name)
+
+        model = await streamer.load_model_from_s3(model_name)
+        tokenizer = await streamer.load_tokenizer_from_s3(model_name)
+
+        if task not in PIPELINE_MAP:
+            raise HTTPException(status_code=400, detail="Pipeline task no soportado")
+
+        nlp_pipeline = pipeline(PIPELINE_MAP[task], model=model, tokenizer=tokenizer)
+
+        result = await asyncio.to_thread(nlp_pipeline, input_text)
+
+        if len(result) > MAX_TOKENS:
+            chunks = split_text_by_tokens(result, tokenizer)
+            full_result = ""
+            for chunk in chunks:
+                full_result += continue_generation(chunk, model, tokenizer)
+            return {"result": full_result}
+        
+        return {"result": result}
+
+    except HTTPException as e:
+        logger.error(f"Error al realizar la predicción: {str(e.detail)}")
+        return JSONResponse(status_code=e.status_code, content={"detail": str(e.detail)})
+
+    except Exception as e:
+        logger.error(f"Error inesperado: {str(e)}")
+        return JSONResponse(status_code=500, content={"detail": "Error inesperado. Intenta más tarde."})
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)