Spaces:

Pixelkram
/

next-token-explorer

Paused

Upload 7 files

f408b78 verified about 1 month ago

12 kB

"""Next Token Explorer: live next-token probabilities from a small language model.
Gradio Space, runs on ZeroGPU or on a free CPU. The whole interface is one gr.HTML component
(ui/template.html, ui/style.css, ui/app.js):
- tokenizing while you type is a server function (CPU only, no queue, no GPU),
- every model step goes through a normal queued Gradio event, so ZeroGPU can attribute GPU time
  to the visitor. One GPU call can take several steps (used by "Write on its own").
"""
import spaces  # import before torch touches CUDA (ZeroGPU)

import json
import os
import time
from pathlib import Path

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-1.5B-Instruct")
DTYPE_NAME = os.environ.get("DTYPE", "float32")
DTYPE = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[DTYPE_NAME]
ZERO_GPU = spaces.config.Config.zero_gpu
DEVICE = "cuda" if (ZERO_GPU or torch.cuda.is_available()) else "cpu"
TOP_K = 20
MAX_INPUT_TOKENS = int(os.environ.get("MAX_INPUT_TOKENS", "1024"))
MAX_STEPS = 40
GRID = [round(0.05 * i, 2) for i in range(1, 41)]  # temperatures 0.05 ... 2.0
UI = Path(__file__).parent / "ui"

if DEVICE == "cpu":
    torch.set_num_threads(max(1, os.cpu_count() or 1))
torch.set_grad_enabled(False)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=DTYPE)
model.to(DEVICE)  # on ZeroGPU this must happen at module level
model.eval()
EOS_IDS = sorted({i for i in [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|im_end|>"),
                              tokenizer.convert_tokens_to_ids("<|endoftext|>")] if i is not None})
VOCAB = len(tokenizer)
_probe = tokenizer.apply_chat_template([{"role": "user", "content": "x"}], tokenize=False)
DEFAULT_SYSTEM = _probe.split("<|im_start|>system\n", 1)[1].split("<|im_end|>", 1)[0] if "<|im_start|>system\n" in _probe else ""
IM_START, IM_END = tokenizer.convert_tokens_to_ids("<|im_start|>"), tokenizer.convert_tokens_to_ids("<|im_end|>")
SYS_HEAD = len(tokenizer("<|im_start|>system\n", add_special_tokens=False).input_ids)   # <|im_start|> system \n
OUTPUT_SIZE = model.config.vocab_size


# ---------------------------------------------------------------- tokens as readable pieces
def _bytes_to_unicode():
    bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
    cs = bs[:]
    n = 0
    for b in range(256):
        if b not in bs:
            bs.append(b)
            cs.append(256 + n)
            n += 1
    return {chr(c): b for b, c in zip(bs, cs)}


U2B = _bytes_to_unicode()
SPECIAL_IDS = set(tokenizer.all_special_ids)
_pieces = {}


def piece(token_id: int) -> dict:
    """Readable form of one token: its text, or its raw bytes if it holds only part of a character."""
    if token_id in _pieces:
        return _pieces[token_id]
    raw = tokenizer.convert_ids_to_tokens(token_id)
    if token_id in SPECIAL_IDS or (raw.startswith("<|") and raw.endswith("|>")):
        out = {"id": token_id, "s": raw, "special": True}
    else:
        try:
            b = bytes(U2B[ch] for ch in raw)
        except KeyError:
            b = raw.encode("utf-8")
        try:
            out = {"id": token_id, "s": b.decode("utf-8")}
        except UnicodeDecodeError:
            out = {"id": token_id, "s": " ".join(f"{x:02x}" for x in b), "partial": True}
    _pieces[token_id] = out
    return out


def chat_messages(text: str, system: str | None) -> list[dict]:
    """An empty system prompt means no system message: Qwen's chat template then inserts its own default."""
    msgs = [{"role": "system", "content": system}] if (system or "").strip() else []
    return msgs + [{"role": "user", "content": text}]


def build_ids(text: str, mode: str, cont: list, system: str | None = None) -> list[int]:
    text = (text or "")[:4000]
    system = (system or "")[:4000] if isinstance(system, str) else ""
    if mode == "chat":
        prompt = tokenizer.apply_chat_template(chat_messages(text, system), tokenize=False, add_generation_prompt=True)
    else:
        prompt = text
    cont = [int(i) for i in (cont or [])][:400]
    if any(i < 0 or i >= VOCAB for i in cont):
        raise ValueError("Unknown token in the continuation.")
    ids = tokenizer(prompt, add_special_tokens=False).input_ids + cont
    if len(ids) > MAX_INPUT_TOKENS:
        raise ValueError(f"The input is too long for this demo ({len(ids)} tokens, limit {MAX_INPUT_TOKENS}). Shorten the text or the system prompt.")
    return ids


# ---------------------------------------------------------------- the model (GPU on ZeroGPU)
_cpu_cache = {"ids": [], "kv": None}  # CPU only: reuse earlier tokens between calls


def _forward(seq, new, past):
    out = model(input_ids=torch.tensor([new], device=DEVICE), past_key_values=past, use_cache=True)
    return torch.log_softmax(out.logits[0, -1].float(), dim=-1), out.past_key_values


def _summary(lp):
    top = torch.topk(lp, TOP_K)
    temps = torch.tensor(GRID, device=lp.device).unsqueeze(1)
    lD = torch.logsumexp(lp.unsqueeze(0) / temps, dim=1)  # log sum_j p_j^(1/T), exact over the whole output
    return {"ids": top.indices.tolist(), "lp": [round(v, 5) for v in top.values.tolist()], "lD": [round(v, 5) for v in lD.tolist()]}


def _duration(ids, actions, temperature):
    # seconds of GPU time to reserve (on ZeroGPU), the actual compute is a small fraction of this
    return min(40, 10 + 0.25 * len(actions))


@spaces.GPU(duration=_duration)
def advance(ids: list[int], actions: list[dict], temperature: float) -> dict:
    """Run the model on ids, then apply each action: {"id": token} adds that token,
    {"r": x} samples the token at position x of the cumulative distribution at this temperature
    (exact, over the whole output layer). Returns one summary per state and the chosen tokens."""
    with torch.inference_mode():
        return _advance(ids, actions, temperature)


def _advance(ids, actions, temperature):
    t0 = time.time()
    seq = list(ids)
    if not ZERO_GPU and _cpu_cache["kv"] is not None and 0 < len(_cpu_cache["ids"]) < len(seq) and seq[:len(_cpu_cache["ids"])] == _cpu_cache["ids"]:
        lp, past = _forward(seq, seq[len(_cpu_cache["ids"]):], _cpu_cache["kv"])
    else:
        lp, past = _forward(seq, seq, None)
    states, choices = [_summary(lp)], []
    for a in actions[:MAX_STEPS]:
        if "id" in a:
            tid = int(a["id"])
            if not 0 <= tid < VOCAB:
                break
            p_t = torch.softmax(lp / temperature, dim=-1) if temperature > 0 else torch.exp(lp)
            rank = int((p_t > p_t[tid]).sum()) + 1
            p = float(p_t[tid])
        elif temperature == 0:
            tid, rank, p = int(torch.argmax(lp)), 1, 1.0
        else:
            probs = torch.softmax(lp / temperature, dim=-1)
            vals, order = torch.sort(probs, descending=True)
            k = int(torch.searchsorted(torch.cumsum(vals, 0), torch.tensor(float(a["r"]), device=vals.device)).clamp(max=len(vals) - 1))
            tid, rank, p = int(order[k]), k + 1, float(vals[k])
            if tid >= VOCAB:  # unused padding slot of the output layer
                tid, rank, p = int(order[0]), 1, float(vals[0])
        choices.append({"id": tid, "rank": rank, "p": p})
        seq.append(tid)
        lp, past = _forward(seq, [tid], past)
        states.append(_summary(lp))
        if tid in EOS_IDS:
            break
    if not ZERO_GPU:
        _cpu_cache["ids"], _cpu_cache["kv"] = seq, past
    return {"states": states, "choices": choices, "ms": round((time.time() - t0) * 1000)}


# ---------------------------------------------------------------- handlers
def system_span(ids: list[int], mode: str):
    """[first, end) token positions of the system prompt text, so the page can colour it. None outside chat mode."""
    if mode != "chat" or not ids or ids[0] != IM_START or IM_END not in ids:
        return None
    end = ids.index(IM_END)
    return [min(SYS_HEAD, end), end]


def tokenize(payload: dict) -> dict:
    """Server function for live tokenization while typing (CPU only, not queued, no GPU).
    Called from the page as server.tokenize({text, mode, cont, system})."""
    payload = payload if isinstance(payload, dict) else {}
    cont = payload.get("cont") or []
    system = payload.get("system") if isinstance(payload.get("system"), str) else ""
    try:
        ids = build_ids(payload.get("text", ""), payload.get("mode", "raw"), cont, system)
    except ValueError as e:
        return {"error": str(e)}
    return {"tokens": [piece(i) for i in ids], "n": len(ids), "n_cont": len(cont),
            "sys": system_span(ids, payload.get("mode", "raw")), "system_default": not system.strip()}


def info(_=None) -> dict:
    return {"model": MODEL_ID, "dtype": DTYPE_NAME, "device": "ZeroGPU" if ZERO_GPU else DEVICE.upper(),
            "default_system": DEFAULT_SYSTEM,
            "vocab": VOCAB, "output_size": OUTPUT_SIZE, "grid": GRID, "eos": EOS_IDS}


def handle(value: str):
    """Queued event: the page writes a request into the component's value and fires `submit`,
    the reply goes back into the value, where the page picks it up."""
    try:
        req = json.loads(value or "{}")
    except json.JSONDecodeError:
        return gr.skip()
    if req.get("kind") != "req":
        return gr.skip()
    rid = req.get("rid")
    try:
        ids = build_ids(req.get("text", ""), req.get("mode", "raw"), req.get("cont", []), req.get("system"))
        if not ids:
            raise ValueError("Type a start text first.")
        temperature = float(req.get("temperature", 1.0))
        if not 0 <= temperature <= 5:
            raise ValueError("Temperature out of range.")
        actions = [a for a in req.get("actions", []) if isinstance(a, dict)][:MAX_STEPS]
        res = advance(ids, actions, temperature)
        states = [{"top": [dict(piece(i), lp=l) for i, l in zip(s["ids"], s["lp"])], "lD": s["lD"]} for s in res["states"]]
        choices = [dict(piece(c["id"]), rank=c["rank"], p=c["p"]) for c in res["choices"]]
        final_ids = ids + [c["id"] for c in res["choices"]]
        reply = {"kind": "res", "rid": rid, "states": states, "choices": choices, "ms": res["ms"],
                 "tokens": [piece(i) for i in final_ids], "n": len(final_ids),
                 "n_cont": len(req.get("cont", [])) + len(choices), "sys": system_span(ids, req.get("mode", "raw"))}
    except Exception as e:  # shown in the page, including ZeroGPU quota messages
        msg = getattr(e, "message", None) or str(e) or type(e).__name__
        reply = {"kind": "res", "rid": rid, "error": msg}
    return json.dumps(reply)


# ---------------------------------------------------------------- page
with gr.Blocks(title="Next Token Explorer") as demo:
    ui = gr.HTML(
        value="",
        html_template=(UI / "template.html").read_text(encoding="utf-8"),
        js_on_load=(UI / "app.js").read_text(encoding="utf-8"),
        apply_default_css=False,
        elem_id="nte",
        server_functions=[tokenize, info],
    )
    # on a CPU one request at a time (shared key/value cache), on ZeroGPU several in parallel
    ui.submit(handle, inputs=ui, outputs=ui, show_progress="hidden", queue=True,
              concurrency_limit=4 if ZERO_GPU else 1)

if __name__ == "__main__":
    demo.launch(
        css=(UI / "style.css").read_text(encoding="utf-8"),
        head='<link rel="preconnect" href="https://fonts.googleapis.com"><link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>'
             '<link rel="stylesheet" href="https://fonts.googleapis.com/css2?family=Atkinson+Hyperlegible+Mono:wght@400;600&family=Atkinson+Hyperlegible+Next:wght@400;600;700&family=Bricolage+Grotesque:opsz,wght@12..96,500..800&display=swap">',
        footer_links=[],
        ssr_mode=False,
    )

	"""Next Token Explorer: live next-token probabilities from a small language model.

	Gradio Space, runs on ZeroGPU or on a free CPU. The whole interface is one gr.HTML component
	(ui/template.html, ui/style.css, ui/app.js):
	- tokenizing while you type is a server function (CPU only, no queue, no GPU),
	- every model step goes through a normal queued Gradio event, so ZeroGPU can attribute GPU time
	to the visitor. One GPU call can take several steps (used by "Write on its own").
	"""
	import spaces # import before torch touches CUDA (ZeroGPU)

	import json
	import os
	import time
	from pathlib import Path

	import gradio as gr
	import torch
	from transformers import AutoModelForCausalLM, AutoTokenizer

	MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-1.5B-Instruct")
	DTYPE_NAME = os.environ.get("DTYPE", "float32")
	DTYPE = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[DTYPE_NAME]
	ZERO_GPU = spaces.config.Config.zero_gpu
	DEVICE = "cuda" if (ZERO_GPU or torch.cuda.is_available()) else "cpu"
	TOP_K = 20
	MAX_INPUT_TOKENS = int(os.environ.get("MAX_INPUT_TOKENS", "1024"))
	MAX_STEPS = 40
	GRID = [round(0.05 * i, 2) for i in range(1, 41)] # temperatures 0.05 ... 2.0
	UI = Path(__file__).parent / "ui"

	if DEVICE == "cpu":
	torch.set_num_threads(max(1, os.cpu_count() or 1))
	torch.set_grad_enabled(False)

	tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
	model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=DTYPE)
	model.to(DEVICE) # on ZeroGPU this must happen at module level
	model.eval()
	EOS_IDS = sorted({i for i in [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<\|im_end\|>"),
	tokenizer.convert_tokens_to_ids("<\|endoftext\|>")] if i is not None})
	VOCAB = len(tokenizer)
	_probe = tokenizer.apply_chat_template([{"role": "user", "content": "x"}], tokenize=False)
	DEFAULT_SYSTEM = _probe.split("<\|im_start\|>system\n", 1)[1].split("<\|im_end\|>", 1)[0] if "<\|im_start\|>system\n" in _probe else ""
	IM_START, IM_END = tokenizer.convert_tokens_to_ids("<\|im_start\|>"), tokenizer.convert_tokens_to_ids("<\|im_end\|>")
	SYS_HEAD = len(tokenizer("<\|im_start\|>system\n", add_special_tokens=False).input_ids) # <\|im_start\|> system \n
	OUTPUT_SIZE = model.config.vocab_size


	# ---------------------------------------------------------------- tokens as readable pieces
	def _bytes_to_unicode():
	bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
	cs = bs[:]
	n = 0
	for b in range(256):
	if b not in bs:
	bs.append(b)
	cs.append(256 + n)
	n += 1
	return {chr(c): b for b, c in zip(bs, cs)}


	U2B = _bytes_to_unicode()
	SPECIAL_IDS = set(tokenizer.all_special_ids)
	_pieces = {}


	def piece(token_id: int) -> dict:
	"""Readable form of one token: its text, or its raw bytes if it holds only part of a character."""
	if token_id in _pieces:
	return _pieces[token_id]
	raw = tokenizer.convert_ids_to_tokens(token_id)
	if token_id in SPECIAL_IDS or (raw.startswith("<\|") and raw.endswith("\|>")):
	out = {"id": token_id, "s": raw, "special": True}
	else:
	try:
	b = bytes(U2B[ch] for ch in raw)
	except KeyError:
	b = raw.encode("utf-8")
	try:
	out = {"id": token_id, "s": b.decode("utf-8")}
	except UnicodeDecodeError:
	out = {"id": token_id, "s": " ".join(f"{x:02x}" for x in b), "partial": True}
	_pieces[token_id] = out
	return out


	def chat_messages(text: str, system: str \| None) -> list[dict]:
	"""An empty system prompt means no system message: Qwen's chat template then inserts its own default."""
	msgs = [{"role": "system", "content": system}] if (system or "").strip() else []
	return msgs + [{"role": "user", "content": text}]


	def build_ids(text: str, mode: str, cont: list, system: str \| None = None) -> list[int]:
	text = (text or "")[:4000]
	system = (system or "")[:4000] if isinstance(system, str) else ""
	if mode == "chat":
	prompt = tokenizer.apply_chat_template(chat_messages(text, system), tokenize=False, add_generation_prompt=True)
	else:
	prompt = text
	cont = [int(i) for i in (cont or [])][:400]
	if any(i < 0 or i >= VOCAB for i in cont):
	raise ValueError("Unknown token in the continuation.")
	ids = tokenizer(prompt, add_special_tokens=False).input_ids + cont
	if len(ids) > MAX_INPUT_TOKENS:
	raise ValueError(f"The input is too long for this demo ({len(ids)} tokens, limit {MAX_INPUT_TOKENS}). Shorten the text or the system prompt.")
	return ids


	# ---------------------------------------------------------------- the model (GPU on ZeroGPU)
	_cpu_cache = {"ids": [], "kv": None} # CPU only: reuse earlier tokens between calls


	def _forward(seq, new, past):
	out = model(input_ids=torch.tensor([new], device=DEVICE), past_key_values=past, use_cache=True)
	return torch.log_softmax(out.logits[0, -1].float(), dim=-1), out.past_key_values


	def _summary(lp):
	top = torch.topk(lp, TOP_K)
	temps = torch.tensor(GRID, device=lp.device).unsqueeze(1)
	lD = torch.logsumexp(lp.unsqueeze(0) / temps, dim=1) # log sum_j p_j^(1/T), exact over the whole output
	return {"ids": top.indices.tolist(), "lp": [round(v, 5) for v in top.values.tolist()], "lD": [round(v, 5) for v in lD.tolist()]}


	def _duration(ids, actions, temperature):
	# seconds of GPU time to reserve (on ZeroGPU), the actual compute is a small fraction of this
	return min(40, 10 + 0.25 * len(actions))


	@spaces.GPU(duration=_duration)
	def advance(ids: list[int], actions: list[dict], temperature: float) -> dict:
	"""Run the model on ids, then apply each action: {"id": token} adds that token,
	{"r": x} samples the token at position x of the cumulative distribution at this temperature
	(exact, over the whole output layer). Returns one summary per state and the chosen tokens."""
	with torch.inference_mode():
	return _advance(ids, actions, temperature)


	def _advance(ids, actions, temperature):
	t0 = time.time()
	seq = list(ids)
	if not ZERO_GPU and _cpu_cache["kv"] is not None and 0 < len(_cpu_cache["ids"]) < len(seq) and seq[:len(_cpu_cache["ids"])] == _cpu_cache["ids"]:
	lp, past = _forward(seq, seq[len(_cpu_cache["ids"]):], _cpu_cache["kv"])
	else:
	lp, past = _forward(seq, seq, None)
	states, choices = [_summary(lp)], []
	for a in actions[:MAX_STEPS]:
	if "id" in a:
	tid = int(a["id"])
	if not 0 <= tid < VOCAB:
	break
	p_t = torch.softmax(lp / temperature, dim=-1) if temperature > 0 else torch.exp(lp)
	rank = int((p_t > p_t[tid]).sum()) + 1
	p = float(p_t[tid])
	elif temperature == 0:
	tid, rank, p = int(torch.argmax(lp)), 1, 1.0
	else:
	probs = torch.softmax(lp / temperature, dim=-1)
	vals, order = torch.sort(probs, descending=True)
	k = int(torch.searchsorted(torch.cumsum(vals, 0), torch.tensor(float(a["r"]), device=vals.device)).clamp(max=len(vals) - 1))
	tid, rank, p = int(order[k]), k + 1, float(vals[k])
	if tid >= VOCAB: # unused padding slot of the output layer
	tid, rank, p = int(order[0]), 1, float(vals[0])
	choices.append({"id": tid, "rank": rank, "p": p})
	seq.append(tid)
	lp, past = _forward(seq, [tid], past)
	states.append(_summary(lp))
	if tid in EOS_IDS:
	break
	if not ZERO_GPU:
	_cpu_cache["ids"], _cpu_cache["kv"] = seq, past
	return {"states": states, "choices": choices, "ms": round((time.time() - t0) * 1000)}


	# ---------------------------------------------------------------- handlers
	def system_span(ids: list[int], mode: str):
	"""[first, end) token positions of the system prompt text, so the page can colour it. None outside chat mode."""
	if mode != "chat" or not ids or ids[0] != IM_START or IM_END not in ids:
	return None
	end = ids.index(IM_END)
	return [min(SYS_HEAD, end), end]


	def tokenize(payload: dict) -> dict:
	"""Server function for live tokenization while typing (CPU only, not queued, no GPU).
	Called from the page as server.tokenize({text, mode, cont, system})."""
	payload = payload if isinstance(payload, dict) else {}
	cont = payload.get("cont") or []
	system = payload.get("system") if isinstance(payload.get("system"), str) else ""
	try:
	ids = build_ids(payload.get("text", ""), payload.get("mode", "raw"), cont, system)
	except ValueError as e:
	return {"error": str(e)}
	return {"tokens": [piece(i) for i in ids], "n": len(ids), "n_cont": len(cont),
	"sys": system_span(ids, payload.get("mode", "raw")), "system_default": not system.strip()}


	def info(_=None) -> dict:
	return {"model": MODEL_ID, "dtype": DTYPE_NAME, "device": "ZeroGPU" if ZERO_GPU else DEVICE.upper(),
	"default_system": DEFAULT_SYSTEM,
	"vocab": VOCAB, "output_size": OUTPUT_SIZE, "grid": GRID, "eos": EOS_IDS}


	def handle(value: str):
	"""Queued event: the page writes a request into the component's value and fires `submit`,
	the reply goes back into the value, where the page picks it up."""
	try:
	req = json.loads(value or "{}")
	except json.JSONDecodeError:
	return gr.skip()
	if req.get("kind") != "req":
	return gr.skip()
	rid = req.get("rid")
	try:
	ids = build_ids(req.get("text", ""), req.get("mode", "raw"), req.get("cont", []), req.get("system"))
	if not ids:
	raise ValueError("Type a start text first.")
	temperature = float(req.get("temperature", 1.0))
	if not 0 <= temperature <= 5:
	raise ValueError("Temperature out of range.")
	actions = [a for a in req.get("actions", []) if isinstance(a, dict)][:MAX_STEPS]
	res = advance(ids, actions, temperature)
	states = [{"top": [dict(piece(i), lp=l) for i, l in zip(s["ids"], s["lp"])], "lD": s["lD"]} for s in res["states"]]
	choices = [dict(piece(c["id"]), rank=c["rank"], p=c["p"]) for c in res["choices"]]
	final_ids = ids + [c["id"] for c in res["choices"]]
	reply = {"kind": "res", "rid": rid, "states": states, "choices": choices, "ms": res["ms"],
	"tokens": [piece(i) for i in final_ids], "n": len(final_ids),
	"n_cont": len(req.get("cont", [])) + len(choices), "sys": system_span(ids, req.get("mode", "raw"))}
	except Exception as e: # shown in the page, including ZeroGPU quota messages
	msg = getattr(e, "message", None) or str(e) or type(e).__name__
	reply = {"kind": "res", "rid": rid, "error": msg}
	return json.dumps(reply)


	# ---------------------------------------------------------------- page
	with gr.Blocks(title="Next Token Explorer") as demo:
	ui = gr.HTML(
	value="",
	html_template=(UI / "template.html").read_text(encoding="utf-8"),
	js_on_load=(UI / "app.js").read_text(encoding="utf-8"),
	apply_default_css=False,
	elem_id="nte",
	server_functions=[tokenize, info],
	)
	# on a CPU one request at a time (shared key/value cache), on ZeroGPU several in parallel
	ui.submit(handle, inputs=ui, outputs=ui, show_progress="hidden", queue=True,
	concurrency_limit=4 if ZERO_GPU else 1)

	if __name__ == "__main__":
	demo.launch(
	css=(UI / "style.css").read_text(encoding="utf-8"),
	head='<link rel="preconnect" href="https://fonts.googleapis.com"><link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>'
	'<link rel="stylesheet" href="https://fonts.googleapis.com/css2?family=Atkinson+Hyperlegible+Mono:wght@400;600&family=Atkinson+Hyperlegible+Next:wght@400;600;700&family=Bricolage+Grotesque:opsz,wght@12..96,500..800&display=swap">',
	footer_links=[],
	ssr_mode=False,
	)