app.py

from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel
import uvicorn
import requests
import io
import asyncio
from typing import List, Dict, Any
from llama_cpp import Llama  # Ajusta según la biblioteca que estés utilizando

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"}
]

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)

    async def download_model_to_memory(self, model_config):
        url = f"https://huggingface.co/{model_config['repo_id']}/resolve/main/{model_config['filename']}"
        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:
            llama = model_data['model']
            tokenizer = model_data['tokenizer']

            response = await asyncio.get_event_loop().run_in_executor(
                None,
                lambda: llama.generate(user_input, max_length=1000, do_sample=True)
            )
            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 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:
        model_manager = ModelManager()
        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)

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()