Update app.py

app.py

@@ -3,12 +3,11 @@ import json
 import logging
 import boto3
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import JSONResponse
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from fastAPI.responses import JSONResponse
+from transformers import AutoModelForCausalLM, AutoTokenizer
 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()
@@ -33,17 +32,6 @@ s3_client = boto3.client(
 
 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(
@@ -73,33 +61,31 @@ class S3DirectStream:
 
     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']]
+            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:
-            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)
+            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_prefix, model)
+                await self.download_and_upload_to_s3(model_name)
 
-            config_stream = await self.stream_from_s3(f"{model_prefix}/config.json")
+            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_prefix}/config.json está vacío o no se pudo leer.")
+                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_prefix}", config=config_json, from_tf=False)
+            model = AutoModelForCausalLM.from_pretrained(f"s3://{self.bucket_name}/{model_name}", config=config_json, from_tf=False)
             return model
 
         except HTTPException as e:
@@ -109,21 +95,20 @@ class S3DirectStream:
 
     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")
+            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}/{profile}/{model}")
+            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('/')
+            folder_keys = s3_key.split('-')
             for i in range(1, len(folder_keys)):
-                folder_key = '/'.join(folder_keys[:i]) + '/'
+                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='')
@@ -138,21 +123,48 @@ class S3DirectStream:
         except self.s3_client.exceptions.ClientError:
             return False
 
-    async def download_and_upload_to_s3(self, model_prefix, model_name):
+    async def download_and_upload_to_s3(self, model_name, force_download=False):
         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 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_prefix}/config.json"):
+                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_prefix}/config.json", Body=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_prefix}/tokenizer.json"):
+            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_prefix}/tokenizer.json", Body=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 descargar o cargar archivos desde Hugging Face a S3: {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)
@@ -165,53 +177,26 @@ def split_text_by_tokens(text, tokenizer, max_tokens=MAX_TOKENS):
 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):]
+        chunks = split_text_by_tokens(input_text, tokenizer, max_tokens)
+        for chunk in chunks:
+            generated_text += model.generate(chunk)
     return generated_text
 
-@app.post("/predict/")
-async def predict(model_request: dict):
+@app.post("/generate")
+async def generate_text(model_name: str, input_text: str):
     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)
+        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)
         
-        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}
+        chunks = split_text_by_tokens(input_text, tokenizer, max_tokens=MAX_TOKENS)
         
-        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)})
+        generated_text = continue_generation(input_text, model, tokenizer)
+        
+        return {"generated_text": generated_text}
 
     except Exception as e:
-        logger.error(f"Error inesperado: {str(e)}")
-        return JSONResponse(status_code=500, content={"detail": "Error inesperado. Intenta más tarde."})
+        return JSONResponse(status_code=500, content={"detail": str(e)})
 
 if __name__ == "__main__":
     import uvicorn