File size: 8,863 Bytes
8c32c92 f8a1e1e 8c32c92 7af609e 8c32c92 f8a1e1e 8c32c92 80e2eea 8c32c92 f8a1e1e 8c32c92 f8a1e1e 8c32c92 f8a1e1e 80e2eea f8a1e1e 0a972b6 f8a1e1e 8c32c92 f8a1e1e 8c32c92 f8a1e1e 8c32c92 f8a1e1e |
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 170 171 172 173 174 175 176 177 178 179 180 181 |
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/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) |