Update app.py

app.py

@@ -1,11 +1,11 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from llama_cpp import Llama
-from concurrent.futures import ProcessPoolExecutor, as_completed
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from tqdm import tqdm
 import uvicorn
 from dotenv import load_dotenv
 from difflib import SequenceMatcher
-from tqdm import tqdm
 import threading
 
 load_dotenv()
@@ -13,30 +13,23 @@ load_dotenv()
 app = FastAPI()
 
 # Configuración de los modelos
-models = [
+model_configs = [
     {"repo_id": "Ffftdtd5dtft/gpt2-xl-Q2_K-GGUF", "filename": "gpt2-xl-q2_k.gguf"},
     {"repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-8B-Instruct-Q2_K-GGUF", "filename": "meta-llama-3.1-8b-instruct-q2_k.gguf"},
     {"repo_id": "Ffftdtd5dtft/gemma-2-9b-it-Q2_K-GGUF", "filename": "gemma-2-9b-it-q2_k.gguf"},
     {"repo_id": "Ffftdtd5dtft/gemma-2-27b-Q2_K-GGUF", "filename": "gemma-2-27b-q2_k.gguf"},
 ]
 
-# Función para cargar un modelo
+# Cargar un modelo
 def load_model(model_config):
     return Llama.from_pretrained(repo_id=model_config['repo_id'], filename=model_config['filename'])
 
-# Cargar modelos en paralelo
+# Cargar todos los modelos simultáneamente
 def load_all_models():
-    with ProcessPoolExecutor() as executor:
-        future_to_model = {executor.submit(load_model, model): model for model in models}
-        loaded_models = {}
-        for future in as_completed(future_to_model):
-            model = future_to_model[future]
-            try:
-                loaded_models[model['repo_id']] = future.result()
-                print(f"Modelo cargado en RAM: {model['repo_id']}")
-            except Exception as exc:
-                print(f"Error al cargar modelo {model['repo_id']}: {exc}")
-    return loaded_models
+    with ThreadPoolExecutor(max_workers=len(model_configs)) as executor:
+        futures = [executor.submit(load_model, config) for config in model_configs]
+        models = [future.result() for future in as_completed(futures)]
+    return models
 
 # Cargar modelos en memoria
 llms = load_all_models()
@@ -47,7 +40,7 @@ class ChatRequest(BaseModel):
     top_p: float = 0.95
     temperature: float = 0.7
 
-# Función global para generar respuestas de chat
+# Función para generar respuestas de chat
 def generate_chat_response(request, llm):
     try:
         user_input = normalize_input(request.message)
@@ -103,8 +96,10 @@ def filter_by_similarity(responses):
             break
     return best_response
 
-def worker_function(llm, request):
-    return generate_chat_response(request, llm)
+def worker_function(llm, request, progress_bar):
+    response = generate_chat_response(request, llm)
+    progress_bar.update(1)
+    return response
 
 @app.post("/generate_chat")
 async def generate_chat(request: ChatRequest):
@@ -114,26 +109,28 @@ async def generate_chat(request: ChatRequest):
     print(f"Procesando solicitud: {request.message}")
 
     responses = []
-    threads = []
-
-    # Crear un hilo para cada modelo
-    for llm in llms.values():
-        thread = threading.Thread(target=lambda: responses.append(worker_function(llm, request)))
-        threads.append(thread)
-        thread.start()
-
-    # Esperar a que todos los hilos terminen
-    for thread in threads:
-        thread.join()
+    num_models = len(llms)
+
+    # Crear barra de progreso
+    with tqdm(total=num_models, desc="Generando respuestas", unit="modelo") as progress_bar:
+        # Ejecutar modelos en paralelo
+        with ThreadPoolExecutor(max_workers=num_models) as executor:
+            futures = [executor.submit(worker_function, llm, request, progress_bar) for llm in llms]
+            for future in as_completed(futures):
+                try:
+                    response = future.result()
+                    responses.append(response['response'])
+                except Exception as exc:
+                    print(f"Error en la generación de respuesta: {exc}")
 
     # Seleccionar la mejor respuesta
-    best_response = select_best_response([response['response'] for response in responses])
+    best_response = select_best_response(responses)
     
     print(f"Mejor respuesta seleccionada: {best_response}")
 
     return {
         "best_response": best_response,
-        "all_responses": [response['response'] for response in responses]
+        "all_responses": responses
     }
 
 if __name__ == "__main__":