Spaces:
Sleeping
Sleeping
import wget | |
import os | |
import asyncio | |
from fastapi import FastAPI | |
from pydantic import BaseModel | |
import executorch | |
import uvicorn | |
# Definir los parámetros de la solicitud utilizando Pydantic | |
class Request(BaseModel): | |
# Valores por defecto para las rutas de los archivos | |
model_path: str = "Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8.pte" | |
tokenizer_path: str = "tokenizer.model" | |
# Valores por defecto para la temperatura y longitud de secuencia | |
prompt: str | |
temperature: float = 0.7 | |
seq_len: int = 256 | |
# Crear una instancia de FastAPI | |
app = FastAPI() | |
# Cargar el modelo y el tokenizer globalmente para evitar recargarlo en cada solicitud | |
model = None | |
tokenizer = None | |
# Función para descargar el modelo y el tokenizer en la misma carpeta del código (ahora es asíncrona) | |
async def download_files(model_url: str, tokenizer_url: str, model_path: str, tokenizer_path: str): | |
# Verificar si el archivo ya existe, si no, descargarlo | |
if not os.path.exists(model_path): | |
print(f"Descargando el modelo desde: {model_url}") | |
# Usamos wget de manera asíncrona | |
await asyncio.to_thread(wget.download, model_url, model_path) | |
if not os.path.exists(tokenizer_path): | |
print(f"Descargando el tokenizer desde: {tokenizer_url}") | |
# Usamos wget de manera asíncrona | |
await asyncio.to_thread(wget.download, tokenizer_url, tokenizer_path) | |
# Cargar el modelo y el tokenizer (se hace una vez al inicio) (ahora es asíncrona) | |
async def load_model(request: Request): | |
global model, tokenizer | |
# Si ya están cargados, no cargamos de nuevo | |
if model is not None and tokenizer is not None: | |
print("Modelo y tokenizer ya están cargados en la memoria.") | |
return | |
# URLs de los archivos a descargar | |
model_url = "https://huggingface.co/executorch-community/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8-ET/resolve/main/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8.pte" | |
tokenizer_url = "https://huggingface.co/executorch-community/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8-ET/resolve/main/tokenizer.model" | |
# Obtener la ruta del directorio actual | |
current_dir = os.path.dirname(os.path.realpath(__file__)) | |
# Definir las rutas locales donde se guardarán los archivos | |
local_model_path = os.path.join(current_dir, request.model_path) | |
local_tokenizer_path = os.path.join(current_dir, request.tokenizer_path) | |
# Descargar los archivos si no existen (se hace de forma asíncrona) | |
await download_files(model_url, tokenizer_url, local_model_path, local_tokenizer_path) | |
# Cargar el modelo y tokenizer descargados | |
print("Cargando el modelo y tokenizer en memoria...") | |
model = executorch.load_model(local_model_path) | |
tokenizer = executorch.load_tokenizer(local_tokenizer_path) | |
# Configurar el modelo con los parámetros recibidos | |
model.set_temperature(request.temperature) | |
model.set_max_length(request.seq_len) | |
print("Modelo y tokenizer cargados en memoria.") | |
# Ruta para generar texto (ahora es asíncrona) | |
async def generate_text(request: Request): | |
# Cargar el modelo si no ha sido cargado | |
if model is None or tokenizer is None: | |
await load_model(request) # Esperamos la carga del modelo de forma asíncrona | |
# Generar el texto con el prompt recibido | |
output = model.generate(request.prompt) | |
# Devolver el texto generado | |
return {"generated_text": output} | |
# Código para ejecutar uvicorn desde el archivo Python | |
if __name__ == "__main__": | |
uvicorn.run(app, host="0.0.0.0", port=7860) | |