Update app.py

app.py

@@ -6,7 +6,6 @@ from fastapi import FastAPI, HTTPException
 from fastapi.responses import JSONResponse
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import asyncio
-import concurrent.futures
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -16,7 +15,7 @@ AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
 AWS_REGION = os.getenv("AWS_REGION")
 S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
 
-MAX_TOKENS = 1024  # Limite de tokens por fragmento
+MAX_TOKENS = 1024
 
 s3_client = boto3.client(
     's3',
@@ -84,7 +83,6 @@ class S3DirectStream:
             if not model_files:
                 raise HTTPException(status_code=404, detail=f"Archivos del modelo {model_name} no encontrados en S3.")
 
-            # Verificar que existe el archivo config.json
             config_stream = await self.stream_from_s3(f"{model_prefix}/config.json")
             config_data = config_stream.read()
 
@@ -148,7 +146,7 @@ def continue_generation(input_text, model, tokenizer, max_tokens=MAX_TOKENS):
         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):]  # Si la entrada se agotó, ya no hay más que procesar
+        input_text = input_text[len(input_text):]
     return generated_text
 
 @app.post("/predict/")
@@ -163,7 +161,7 @@ async def predict(model_request: dict):
 
         streamer = S3DirectStream(S3_BUCKET_NAME)
         
-        await streamer.create_s3_folders(model_name)  # Crear las carpetas si no existen
+        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)