Update app.py

app.py

@@ -2,12 +2,13 @@ from fastapi import FastAPI, HTTPException, Request
 from pydantic import BaseModel
 import uvicorn
 import requests
-import io
 import asyncio
+import os
+import io
+import time
 from typing import List, Dict, Any
 from llama_cpp import Llama  # Ajusta según la biblioteca que estés utilizando
-import os
-
+from tqdm import tqdm
 
 app = FastAPI()
 
@@ -38,96 +39,124 @@ class ModelManager:
         self.load_lock = asyncio.Lock()
         self.index_lock = asyncio.Lock()
         self.part_size = 1024 * 1024  # Tamaño de cada parte en bytes (1 MB)
+        self.max_loading_time = 0  # Tiempo máximo en segundos para cargar un modelo
 
     async def download_model_to_memory(self, model_config):
         url = f"https://huggingface.co/{model_config['repo_id']}/resolve/main/{model_config['filename']}"
+        print(f"Descargando modelo desde {url}")
         try:
             response = requests.get(url)
             response.raise_for_status()
+            print(f"Descarga completa para {model_config['name']}")
             return io.BytesIO(response.content)
         except requests.RequestException as e:
             raise HTTPException(status_code=500, detail=f"Error al descargar el modelo: {e}")
 
+    async def save_model_to_temp_file(self, model_config):
+        model_file = await self.download_model_to_memory(model_config)
+        temp_filename = f"/tmp/{model_config['filename']}"
+        print(f"Guardando el modelo en {temp_filename}")
+        with open(temp_filename, 'wb') as f:
+            f.write(model_file.getvalue())
+        print(f"Modelo guardado en {temp_filename}")
+        return temp_filename
+
     async def load_model(self, model_config):
         async with self.load_lock:
             try:
-                model_file = await self.download_model_to_memory(model_config)
-                llama = Llama(model_file)  # Ajusta según la biblioteca y clase correctas
-
-                tokenizer = llama.tokenizer
-                model_data = {
-                    'model': llama,
-                    'tokenizer': tokenizer,
-                    'pad_token': tokenizer.pad_token,
-                    'pad_token_id': tokenizer.pad_token_id,
-                    'eos_token': tokenizer.eos_token,
-                    'eos_token_id': tokenizer.eos_token_id,
-                    'bos_token': tokenizer.bos_token,
-                    'bos_token_id': tokenizer.bos_token_id,
-                    'unk_token': tokenizer.unk_token,
-                    'unk_token_id': tokenizer.unk_token_id
-                }
-
-                self.models[model_config['name']] = model_data
-                await self.handle_large_model(model_config, model_file)
+                start_time = time.time()
+                temp_filename = await self.save_model_to_temp_file(model_config)
+                elapsed_time = time.time() - start_time
+                if elapsed_time > self.max_loading_time:
+                    print(f"El modelo {model_config['name']} tardó {elapsed_time:.2f} segundos en cargar. Dividiendo el modelo.")
+                    await self.handle_large_model(temp_filename, model_config)
+                else:
+                    print(f"Cargando modelo desde {temp_filename}")
+                    llama = Llama(temp_filename)  # Ajusta según la biblioteca y clase correctas
+
+                    tokenizer = llama.tokenizer
+                    model_data = {
+                        'model': llama,
+                        'tokenizer': tokenizer,
+                        'pad_token': tokenizer.pad_token,
+                        'pad_token_id': tokenizer.pad_token_id,
+                        'eos_token': tokenizer.eos_token,
+                        'eos_token_id': tokenizer.eos_token_id,
+                        'bos_token': tokenizer.bos_token,
+                        'bos_token_id': tokenizer.bos_token_id,
+                        'unk_token': tokenizer.unk_token,
+                        'unk_token_id': tokenizer.unk_token_id
+                    }
+
+                    self.models[model_config['name']] = model_data
+                    print(f"Modelo {model_config['name']} cargado correctamente")
             except Exception as e:
                 print(f"Error al cargar el modelo: {e}")
 
-    async def handle_large_model(self, model_config, model_file):
-        total_size = len(model_file.getvalue())
+    async def handle_large_model(self, model_filename, model_config):
+        total_size = os.path.getsize(model_filename)
         num_parts = (total_size + self.part_size - 1) // self.part_size
 
-        for i in range(num_parts):
-            start = i * self.part_size
-            end = min(start + self.part_size, total_size)
-            model_part = io.BytesIO(model_file.getvalue()[start:end])
-            await self.index_model_part(model_part, i)
+        print(f"Modelo {model_config['name']} dividido en {num_parts} partes")
+        with open(model_filename, 'rb') as file:
+            for i in tqdm(range(num_parts), desc=f"Indexando {model_config['name']}"):
+                start = i * self.part_size
+                end = min(start + self.part_size, total_size)
+                file.seek(start)
+                model_part = io.BytesIO(file.read(end - start))
+                await self.index_model_part(model_part, i)
 
     async def index_model_part(self, model_part, part_index):
         async with self.index_lock:
             part_name = f"part_{part_index}"
+            print(f"Indexando parte {part_index}")
             llama_part = Llama(model_part)
             self.model_parts[part_name] = llama_part
+            print(f"Parte {part_index} indexada")
 
     async def generate_response(self, user_input):
+        print("Generando respuestas")
         tasks = [self.generate_chat_response(user_input, model_data) for model_data in self.models.values()]
         responses = await asyncio.gather(*tasks)
         return responses
 
     async def generate_chat_response(self, user_input, model_data):
         try:
-            llama = model_data['model']
-            tokenizer = model_data['tokenizer']
-
-            response = await asyncio.get_event_loop().run_in_executor(
-                None,
-                lambda: llama.generate(user_input, max_length=1000, do_sample=True)
-            )
-            generated_text = response['generated_text']
-
-            # Dividir el texto generado en partes si es necesario
-            parts = []
-            while len(generated_text) > 1000:
-                part = generated_text[:1000]
-                generated_text = generated_text[1000:]
-                parts.append(part.strip())
-            if generated_text:
-                parts.append(generated_text.strip())
-
-            return {"response": '\n'.join(parts), "model_name": model_data['name']}
+            print(f"Generando respuesta usando el modelo {model_data['model']}")
+            start_time = time.time()
+            generated_text = model_data['model'].generate(user_input)
+            elapsed_time = time.time() - start_time
+
+            if len(generated_text) > 1000:
+                parts = []
+                while len(generated_text) > 1000:
+                    part = generated_text[:1000]
+                    parts.append(part)
+                    generated_text = generated_text[1000:]
+                parts.append(generated_text)
+            else:
+                parts = [generated_text]
+
+            print(f"Respuesta generada usando el modelo {model_data['model']} en {elapsed_time:.2f} segundos")
+            return {
+                'model_name': model_data['model'],
+                'generated_text_parts': parts
+            }
         except Exception as e:
-            print(f"Error al generar la respuesta: {e}")
-            return {"response": "Error al generar la respuesta", "model_name": model_data['name']}
+            print(f"Error al generar respuesta con el modelo {model_data['model']}: {e}")
+            return {'model_name': model_data['model'], 'error': str(e)}
 
-@app.post("/chat")
-async def chat(request: Request):
-    body = await request.json()
-    user_input = body.get('message', '').strip()
+@app.post("/generate/")
+async def generate(request: Request):
+    data = await request.json()
+    user_input = data.get('input', '')
     if not user_input:
-        raise HTTPException(status_code=400, detail="El mensaje no puede estar vacío.")
+        raise HTTPException(status_code=400, detail="Se requiere una entrada de usuario.")
 
     try:
         model_manager = ModelManager()
+        tasks = [model_manager.load_model(config) for config in model_configs]
+        await asyncio.gather(*tasks)
         responses = await model_manager.generate_response(user_input)
         return {"responses": responses}
     except Exception as e: