File size: 8,456 Bytes
3e1e0dc
9883ddb
 
3e1e0dc
 
9883ddb
47a886a
3e1e0dc
 
8a6af4d
9883ddb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3e1e0dc
9883ddb
 
 
 
 
 
3e1e0dc
 
 
 
 
9883ddb
 
 
3e1e0dc
9883ddb
 
3e1e0dc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9883ddb
3e1e0dc
 
 
 
 
 
4f0da89
3e1e0dc
 
8a6af4d
3e1e0dc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4f0da89
 
1ff2b2e
 
3e1e0dc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9883ddb
 
3e1e0dc
 
 
 
 
9883ddb
3e1e0dc
 
 
7691684
3e1e0dc
 
8c27558
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
from fastapi import FastAPI, HTTPException, Request
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

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']}"
        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}")
        with open(temp_filename, 'wb') as f:
            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:
            try:
                temp_filename = await self.save_model_to_temp_file(model_config)
                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.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}")

    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 = 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):
        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)

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

if __name__ == "__main__":
    asyncio.run(start_uvicorn())