Update app.py

app.py

@@ -1,19 +1,12 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Request
 from pydantic import BaseModel
-from llama_cpp import Llama
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from tqdm import tqdm
 import uvicorn
-from dotenv import load_dotenv
-import io
 import requests
+import io
 import asyncio
-import time
-
-# Cargar variables de entorno
-load_dotenv()
+from typing import List, Dict, Any
+from llama_cpp import Llama  # Ajusta según la biblioteca que estés utilizando
 
-# Inicializar aplicación FastAPI
 app = FastAPI()
 
 # Configuración de los modelos
@@ -31,149 +24,119 @@ model_configs = [
     {"repo_id": "Ffftdtd5dtft/gemma-2-2b-it-Q2_K-GGUF", "filename": "gemma-2-2b-it-q2_k.gguf", "name": "Gemma 2-2B IT"},
     {"repo_id": "Ffftdtd5dtft/sarvam-2b-v0.5-Q2_K-GGUF", "filename": "sarvam-2b-v0.5-q2_k.gguf", "name": "Sarvam 2B v0.5"},
     {"repo_id": "Ffftdtd5dtft/WizardLM-13B-Uncensored-Q2_K-GGUF", "filename": "wizardlm-13b-uncensored-q2_k.gguf", "name": "WizardLM 13B Uncensored"},
+    {"repo_id": "Ffftdtd5dtft/Qwen2-Math-72B-Instruct-Q2_K-GGUF", "filename": "qwen2-math-72b-instruct-q2_k.gguf", "name": "Qwen2 Math 72B Instruct"},
     {"repo_id": "Ffftdtd5dtft/WizardLM-7B-Uncensored-Q2_K-GGUF", "filename": "wizardlm-7b-uncensored-q2_k.gguf", "name": "WizardLM 7B Uncensored"},
     {"repo_id": "Ffftdtd5dtft/Qwen2-Math-7B-Instruct-Q2_K-GGUF", "filename": "qwen2-math-7b-instruct-q2_k.gguf", "name": "Qwen2 Math 7B Instruct"}
 ]
 
-# Clase para gestionar modelos
 class ModelManager:
     def __init__(self):
-        self.models = []
-        self.configs = {}
+        self.models = {}
+        self.model_parts = {}
+        self.load_lock = asyncio.Lock()
+        self.index_lock = asyncio.Lock()
+        self.part_size = 1024 * 1024  # Tamaño de cada parte en bytes (1 MB)
 
     async def download_model_to_memory(self, model_config):
-        print(f"Descargando modelo: {model_config['name']}...")
         url = f"https://huggingface.co/{model_config['repo_id']}/resolve/main/{model_config['filename']}"
-        response = requests.get(url)
-        if response.status_code == 200:
-            model_file = io.BytesIO(response.content)
-            return model_file
-        else:
-            raise Exception(f"Error al descargar el modelo: {response.status_code}")
+        try:
+            response = requests.get(url)
+            response.raise_for_status()
+            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 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)
+            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())
+        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)
+
+    async def index_model_part(self, model_part, part_index):
+        async with self.index_lock:
+            part_name = f"part_{part_index}"
+            llama_part = Llama(model_part)
+            self.model_parts[part_name] = llama_part
+
+    async def generate_response(self, user_input):
+        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:
-            start_time = time.time()
-            model_file = await self.download_model_to_memory(model_config)
-            print(f"Cargando modelo: {model_config['name']}...")
-            
-            # Cargar el modelo usando llama_cpp
-            llama = await asyncio.get_event_loop().run_in_executor(
+            llama = model_data['model']
+            tokenizer = model_data['tokenizer']
+
+            response = await asyncio.get_event_loop().run_in_executor(
                 None,
-                lambda: Llama.from_pretrained(model_file)
+                lambda: llama.generate(user_input, max_length=1000, do_sample=True)
             )
-            
-            # Simulación de división de carga si el tiempo excede 1 segundo
-            async def load_part(part):
-                # Esta función simula la carga de una parte del modelo
-                await asyncio.sleep(0.1)  # Simula un pequeño retraso en la carga
-
-            if time.time() - start_time > 1:
-                print(f"Modelo {model_config['name']} tardó más de 1 segundo en cargarse, dividiendo la carga...")
-                await asyncio.gather(*(load_part(part) for part in range(5)))  # Simulación de división en 5 partes
-            
-            tokenizer = llama.tokenizer
-
-            # Almacenar tokens y tokenizer en la RAM
-            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.append({"model_data": model_data, "name": model_config['name']})
+            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']}
         except Exception as e:
-            print(f"Error al cargar el modelo: {e}")
-
-    async def load_all_models(self):
-        print("Iniciando carga de modelos...")
-        start_time = time.time()
-        tasks = [self.load_model(config) for config in model_configs]
-        await asyncio.gather(*tasks)
-        end_time = time.time()
-        print(f"Todos los modelos han sido cargados en {end_time - start_time:.2f} segundos.")
-
-# Instanciar ModelManager y cargar modelos
-model_manager = ModelManager()
-
-@app.on_event("startup")
-async def startup_event():
-    await model_manager.load_all_models()
-
-# Modelo global para la solicitud de chat
-class ChatRequest(BaseModel):
-    message: str
-    top_k: int = 50
-    top_p: float = 0.95
-    temperature: float = 0.7
-
-# Límite de tokens para respuestas
-TOKEN_LIMIT = 1000  # Define el límite de tokens permitido por respuesta
-
-# Función para generar respuestas de chat
-async def generate_chat_response(request, model_data):
+            print(f"Error al generar la respuesta: {e}")
+            return {"response": "Error al generar la respuesta", "model_name": model_data['name']}
+
+@app.post("/chat")
+async def chat(request: Request):
+    body = await request.json()
+    user_input = body.get('message', '').strip()
+    if not user_input:
+        raise HTTPException(status_code=400, detail="El mensaje no puede estar vacío.")
+
     try:
-        user_input = normalize_input(request.message)
-        llama = model_data['model_data']['model']
-        tokenizer = model_data['model_data']['tokenizer']
-        
-        # Generar respuesta de manera rápida
-        response = await asyncio.get_event_loop().run_in_executor(
-            None,
-            lambda: llama.generate(user_input, max_length=TOKEN_LIMIT, do_sample=True, top_k=request.top_k, top_p=request.top_p, temperature=request.temperature)
-        )
-        generated_text = response['generated_text']
-        # Dividir respuesta larga
-        split_response = split_long_response(generated_text)
-        return {"response": split_response, "literal": user_input, "model_name": model_data['name']}
+        model_manager = ModelManager()
+        responses = await model_manager.generate_response(user_input)
+        return {"responses": responses}
     except Exception as e:
-        print(f"Error al generar la respuesta: {e}")
-        return {"response": "Error al generar la respuesta", "literal": user_input, "model_name": model_data['name']}
-
-def split_long_response(response):
-    """ Divide la respuesta en partes más pequeñas si excede el límite de tokens. """
-    parts = []
-    while len(response) > TOKEN_LIMIT:
-        part = response[:TOKEN_LIMIT]
-        response = response[TOKEN_LIMIT:]
-        parts.append(part.strip())
-    if response:
-        parts.append(response.strip())
-    return '\n'.join(parts)
-
-def remove_duplicates(text):
-    """ Elimina duplicados en el texto. """
-    lines = text.splitlines()
-    unique_lines = list(dict.fromkeys(lines))
-    return '\n'.join(unique_lines)
-
-def remove_repetitive_responses(responses):
-    unique_responses = []
-    seen_responses = set()
-    for response in responses:
-        normalized_response = remove_duplicates(response['response'])
-        if normalized_response not in seen_responses:
-            seen_responses.add(normalized_response)
-            response['response'] = normalized_response
-            unique_responses.append(response)
-    return unique_responses
+        raise HTTPException(status_code=500, detail=str(e))
 
-@app.post("/chat")
-async def chat(request: ChatRequest):
-    results = []
-    for model_data in model_manager.models:
-        response = await generate_chat_response(request, model_data)
-        results.append(response)
-    unique_results = remove_repetitive_responses(results)
-    return {"results": unique_results}
-
-# Ejecutar la aplicación FastAPI
-if __name__ == "__main__":
+def start_uvicorn():
     uvicorn.run(app, host="0.0.0.0", port=7860)
+
+if __name__ == "__main__":
+    loop = asyncio.get_event_loop()
+    model_manager = ModelManager()
+    tasks = [model_manager.load_model(config) for config in model_configs]
+    loop.run_until_complete(asyncio.gather(*tasks))
+    start_uvicorn()