Update app.py

app.py

@@ -3,35 +3,38 @@ import uvicorn
 import requests
 import os
 import io
-import time
 import asyncio
 from typing import List, Dict, Any
 from tqdm import tqdm
 from llama_cpp import Llama
 import aiofiles
+import time
 
 app = FastAPI()
 
 # Configuración de los 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"},
-    {"repo_id": "Ffftdtd5dtft/gemma-2-9b-it-Q2_K-GGUF", "filename": "gemma-2-9b-it-q2_k.gguf", "name": "Gemma 2-9B IT"},
-    {"repo_id": "Ffftdtd5dtft/gemma-2-27b-Q2_K-GGUF", "filename": "gemma-2-27b-q2_k.gguf", "name": "Gemma 2-27B"},
-    {"repo_id": "Ffftdtd5dtft/Phi-3-mini-128k-instruct-Q2_K-GGUF", "filename": "phi-3-mini-128k-instruct-q2_k.gguf", "name": "Phi-3 Mini 128K Instruct"},
-    {"repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-8B-Q2_K-GGUF", "filename": "meta-llama-3.1-8b-q2_k.gguf", "name": "Meta Llama 3.1-8B"},
-    {"repo_id": "Ffftdtd5dtft/Qwen2-7B-Instruct-Q2_K-GGUF", "filename": "qwen2-7b-instruct-q2_k.gguf", "name": "Qwen2 7B Instruct"},
-    {"repo_id": "Ffftdtd5dtft/starcoder2-3b-Q2_K-GGUF", "filename": "starcoder2-3b-q2_k.gguf", "name": "Starcoder2 3B"},
-    {"repo_id": "Ffftdtd5dtft/Qwen2-1.5B-Instruct-Q2_K-GGUF", "filename": "qwen2-1.5b-instruct-q2_k.gguf", "name": "Qwen2 1.5B Instruct"},
-    {"repo_id": "Ffftdtd5dtft/starcoder2-15b-Q2_K-GGUF", "filename": "starcoder2-15b-q2_k.gguf", "name": "Starcoder2 15B"},
-    {"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"}
+    {"repo_id": "TheBloke/gpt2-xl-Q2_K-GGUF", "filename": "gpt2-xl-q2_k.gguf", "name": "GPT-2 XL"},
+    {"repo_id": "TheBloke/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"},
+    {"repo_id": "TheBloke/gemma-2-9b-it-Q2_K-GGUF", "filename": "gemma-2-9b-it-q2_k.gguf", "name": "Gemma 2-9B IT"},
+    {"repo_id": "TheBloke/gemma-2-27b-Q2_K-GGUF", "filename": "gemma-2-27b-q2_k.gguf", "name": "Gemma 2-27B"},
+    {"repo_id": "TheBloke/Phi-3-mini-128k-instruct-Q2_K-GGUF", "filename": "phi-3-mini-128k-instruct-q2_k.gguf", "name": "Phi-3 Mini 128K Instruct"},
+    {"repo_id": "TheBloke/Meta-Llama-3.1-8B-Q2_K-GGUF", "filename": "meta-llama-3.1-8b-q2_k.gguf", "name": "Meta Llama 3.1-8B"},
+    {"repo_id": "TheBloke/Qwen2-7B-Instruct-Q2_K-GGUF", "filename": "qwen2-7b-instruct-q2_k.gguf", "name": "Qwen2 7B Instruct"},
+    {"repo_id": "TheBloke/starcoder2-3b-Q2_K-GGUF", "filename": "starcoder2-3b-q2_k.gguf", "name": "Starcoder2 3B"},
+    {"repo_id": "TheBloke/Qwen2-1.5B-Instruct-Q2_K-GGUF", "filename": "qwen2-1.5b-instruct-q2_k.gguf", "name": "Qwen2 1.5B Instruct"},
+    {"repo_id": "TheBloke/starcoder2-15b-Q2_K-GGUF", "filename": "starcoder2-15b-q2_k.gguf", "name": "Starcoder2 15B"},
+    {"repo_id": "TheBloke/gemma-2-2b-it-Q2_K-GGUF", "filename": "gemma-2-2b-it-q2_k.gguf", "name": "Gemma 2-2B IT"},
+    {"repo_id": "TheBloke/sarvam-2b-v0.5-Q2_K-GGUF", "filename": "sarvam-2b-v0.5-q2_k.gguf", "name": "Sarvam 2B v0.5"},
+    {"repo_id": "TheBloke/WizardLM-13B-Uncensored-Q2_K-GGUF", "filename": "wizardlm-13b-uncensored-q2_k.gguf", "name": "WizardLM 13B Uncensored"},
+    {"repo_id": "TheBloke/Qwen2-Math-72B-Instruct-Q2_K-GGUF", "filename": "qwen2-math-72b-instruct-q2_k.gguf", "name": "Qwen2 Math 72B Instruct"},
+    {"repo_id": "TheBloke/WizardLM-7B-Uncensored-Q2_K-GGUF", "filename": "wizardlm-7b-uncensored-q2_k.gguf", "name": "WizardLM 7B Uncensored"},
+    {"repo_id": "TheBloke/Qwen2-Math-7B-Instruct-Q2_K-GGUF", "filename": "qwen2-math-7b-instruct-q2_k.gguf", "name": "Qwen2 Math 7B Instruct"}
 ]
 
+# Directorio para almacenar los modelos descargados
+models_dir = "modelos"
+
 class ModelManager:
     def __init__(self):
         self.models = {}
@@ -40,64 +43,66 @@ class ModelManager:
         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):
-        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()
-            end_time = time.time()
-            download_duration = end_time - start_time
-            print(f"Descarga completa para {model_config['name']} en {download_duration:.2f} segundos")
-            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}")
-        async with aiofiles.open(temp_filename, 'wb') as f:
-            await 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:
+    async def download_model(self, model_config):
+        model_path = os.path.join(models_dir, model_config['filename'])
+        if not os.path.exists(model_path):
+            url = f"https://huggingface.co/{model_config['repo_id']}/resolve/main/{model_config['filename']}"
+            print(f"Descargando modelo desde {url}")
             try:
-                temp_filename = await self.save_model_to_temp_file(model_config)
                 start_time = time.time()
-                print(f"Cargando modelo desde {temp_filename}")
-
-                # Cambiar la forma en que se carga el modelo según la biblioteca que utilices
-                llama = Llama.from_file(temp_filename)
-                
+                response = requests.get(url, stream=True)
+                response.raise_for_status()
+
+                total_size = int(response.headers.get('content-length', 0))
+                with open(model_path, 'wb') as f:
+                    with tqdm(total=total_size, unit='B', unit_scale=True, desc=f"Descargando {model_config['filename']}") as pbar:
+                        for chunk in response.iter_content(chunk_size=8192):
+                            f.write(chunk)
+                            pbar.update(len(chunk))
                 end_time = time.time()
-                load_duration = end_time - start_time
-                if load_duration > 0.5:
-                    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}")
+                download_duration = end_time - start_time
+                print(f"Descarga completa para {model_config['name']} en {download_duration:.2f} segundos")
+            except requests.RequestException as e:
+                raise HTTPException(status_code=500, detail=f"Error al descargar el modelo: {e}")
+        else:
+            print(f"Modelo {model_config['filename']} ya descargado.")
+        return model_path
+
+    async def load_model(self, model_config):
+        async with self.load_lock:
+            if model_config['name'] not in self.models:
+                try:
+                    model_path = await self.download_model(model_config)
+                    start_time = time.time()
+                    print(f"Cargando modelo desde {model_path}")
+
+                    llama = Llama(model_path=model_path)
+
+                    end_time = time.time()
+                    load_duration = end_time - start_time
+                    if load_duration > 0.5:
+                        print(f"Modelo {model_config['name']} tardó {load_duration:.2f} segundos en cargar, dividiendo automáticamente")
+                        await self.handle_large_model(model_path, 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)
@@ -116,7 +121,7 @@ class ModelManager:
         async with self.index_lock:
             part_name = f"part_{part_index}"
             print(f"Indexando parte {part_index}")
-            temp_filename = f"/tmp/{part_name}.gguf"
+            temp_filename = os.path.join(models_dir, f"{part_name}.gguf")
             async with aiofiles.open(temp_filename, 'wb') as f:
                 await f.write(model_part.getvalue())
             print(f"Parte {part_index} indexada y guardada")
@@ -156,16 +161,25 @@ async def generate(request: Request):
         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))
 
-def start_uvicorn():
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+async def load_models_on_startup():
+    tasks = [model_manager.load_model(config) for config in model_configs]
+    await asyncio.gather(*tasks)
+
+@app.on_event("startup")
+async def startup_event():
+    global model_manager
+    model_manager = ModelManager()
+    await load_models_on_startup()
+    print("Modelos cargados correctamente. API lista.")
 
 if __name__ == "__main__":
-    asyncio.run(start_uvicorn())
+    # Crear el directorio "modelos" si no existe
+    if not os.path.exists(models_dir):
+        os.makedirs(models_dir)
+
+    uvicorn.run(app, host="0.0.0.0", port=7860)