app.py

from fastapi import FastAPI, HTTPException, Request
import uvicorn
import requests
import os
import io
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"}
]

# Directorio para almacenar los modelos descargados
models_dir = "modelos"

class ModelManager:
    def __init__(self):
        self.models = {}
        self.model_parts = {}
        self.load_lock = asyncio.Lock()
        self.index_lock = asyncio.Lock()
        self.part_size = 102 * 102  # Tamaño de cada parte en bytes (1 MB)

    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:
                start_time = time.time()
                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()
                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:
                        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_id': tokenizer.pad_token_id,
                        'eos_token_id': tokenizer.eos_token_id,
                        'bos_token_id': tokenizer.bos_token_id,
                        'unk_token_id': tokenizer.unk_token_id,
                        'padding_token_id': tokenizer.padding_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']}"):
                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}")
            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")

    async def generate_response(self, user_input):
        results = []
        for model_name, model_data in self.models.items():
            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)

                # Dividir el texto generado en partes
                parts = []
                while len(generated_text) > 1000:
                    part = generated_text[:1000]
                    parts.append(part)
                    generated_text = generated_text[1000:]
                parts.append(generated_text)

                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):
    data = await request.json()
    user_input = data.get('input', '')
    if not user_input:
        raise HTTPException(status_code=400, detail="Se requiere una entrada de usuario.")
    
    try:
        responses = await model_manager.generate_response(user_input)
        return {"responses": responses}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

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__":
    # 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)