Update app.py

app.py

@@ -1,18 +1,16 @@
-from fastapi import FastAPI, HTTPException, Request
-from pydantic import BaseModel
-import uvicorn
-import requests
-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 asyncio
+import requests
 from tqdm import tqdm
+from fastapi import FastAPI, HTTPException, Request
+import uvicorn
+from llama_cpp import Llama
 
 app = FastAPI()
 
-# Configuración de los modelos
+# Configuración de modelos
 model_configs = [
     {"repo_id": "Ffftdtd5dtft/gpt2-xl-Q2_K-GGUF", "filename": "gpt2-xl-q2_k.gguf", "name": "GPT-2 XL"},
     {"repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-8B-Instruct-Q2_K-GGUF", "filename": "meta-llama-3.1-8b-instruct-q2_k.gguf", "name": "Meta Llama 3.1-8B Instruct"},
@@ -35,25 +33,24 @@ model_configs = [
 class ModelManager:
     def __init__(self):
         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)
-        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:
+            start_time = time.time()
             response = requests.get(url)
             response.raise_for_status()
-            print(f"Descarga completa para {model_config['name']}")
-            return io.BytesIO(response.content)
+            model_file = io.BytesIO(response.content)
+            end_time = time.time()
+            download_duration = end_time - start_time
+            print(f"Descarga completa para {model_config['name']} en {download_duration:.2f} segundos")
+            return model_file
         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)
+    async def save_model_to_temp_file(self, model_file, model_config):
         temp_filename = f"/tmp/{model_config['filename']}"
         print(f"Guardando el modelo en {temp_filename}")
         with open(temp_filename, 'wb') as f:
@@ -62,41 +59,39 @@ class ModelManager:
         return temp_filename
 
     async def load_model(self, model_config):
-        async with self.load_lock:
-            try:
-                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}")
+        model_file = await self.download_model_to_memory(model_config)
+        temp_filename = await self.save_model_to_temp_file(model_file, model_config)
+        try:
+            start_time = time.time()
+            print(f"Cargando modelo desde {temp_filename}")
+            llama = Llama.load(temp_filename)
+            end_time = time.time()
+            load_duration = end_time - start_time
+            if load_duration > 0:
+                print(f"Modelo {model_config['name']} tardó {load_duration:.2f} segundos en cargar, dividiendo automáticamente")
+                await self.handle_large_model(temp_filename, model_config)
+            else:
+                print(f"Modelo {model_config['name']} cargado correctamente en {load_duration:.2f} segundos")
+            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
+        except Exception as e:
+            print(f"Error al cargar el modelo: {e}")
 
     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
-
         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']}"):
@@ -107,42 +102,31 @@ class ModelManager:
                 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.from_bytes(model_part.getvalue())  # Ajusta según la biblioteca y clase correctas
-            self.model_parts[part_name] = llama_part
+        part_name = f"part_{part_index}"
+        print(f"Indexando parte {part_index}")
+        temp_filename = f"/tmp/{part_name}.gguf"
+        with open(temp_filename, 'wb') as f:
+            f.write(model_part.getvalue())
+        print(f"Parte {part_index} indexada y guardada")
 
     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:
-            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 respuesta con el modelo {model_data['model']}: {e}")
-            return {'model_name': model_data['model'], 'error': str(e)}
+        results = []
+        for model_name, model_data in self.models.items():
+            print(f"Generando respuesta con el modelo {model_name}")
+            try:
+                tokenizer = model_data['tokenizer']
+                input_ids = tokenizer(user_input, return_tensors="pt").input_ids
+                outputs = model_data['model'].generate(input_ids)
+                generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+                parts = [generated_text[i:i + 1000] for i in range(0, len(generated_text), 1000)]
+                results.append({
+                    'model_name': model_name,
+                    'generated_text_parts': parts
+                })
+            except Exception as e:
+                print(f"Error al generar respuesta con el modelo {model_name}: {e}")
+                results.append({'model_name': model_name, 'error': str(e)})
+        return results
 
 @app.post("/generate/")
 async def generate(request: Request):
@@ -151,21 +135,14 @@ async def generate(request: Request):
     if not user_input:
         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:
-        raise HTTPException(status_code=500, detail=str(e))
+    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}
 
 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()
+    asyncio.run(start_uvicorn())