| """Confidence-guided subword denoising and blockwise continuation.""" |
| import math |
| from pathlib import Path |
| import time |
| import torch |
| from tokenizers import Tokenizer |
| from .model import Denoiser, PAD, BOS, EOS, MASK, UNK |
| from .train import amp |
|
|
|
|
| def load(directory, device="cpu"): |
| return Denoiser.from_pretrained(directory, device), Tokenizer.from_file(str(Path(directory) / "tokenizer.json")) |
|
|
|
|
| @torch.inference_mode() |
| def denoise(model, tokenizer, ids, *, steps=16, temperature=0.7, seed=42, top_k=40, precision="fp32", policy="confidence", threshold=0.85, allow_eos=True): |
| if not ids or len(ids) > model.config.max_length: |
| raise ValueError(f"Input exceeds {model.config.max_length} tokens") |
| if steps < 1 or not math.isfinite(temperature) or temperature < 0 or top_k < 1: |
| raise ValueError("Invalid sampling parameters") |
| if policy not in ("confidence", "left_to_right", "adaptive") or not 0 < threshold <= 1: |
| raise ValueError("Invalid policy or threshold") |
| if not 0 <= seed < 2**63: |
| raise ValueError("seed must be in 0..2^63-1") |
| device = str(next(model.parameters()).device).split(":")[0] |
| generator = torch.Generator(device=device).manual_seed(seed) |
| tokens = torch.tensor([ids], device=device) |
| fixed = tokens.ne(MASK) |
| count = int((~fixed).sum()) |
| iterations = min(steps, count) |
| trace = [] |
| model.eval() |
| started = time.perf_counter() |
| for step in range(iterations): |
| unknown = tokens.eq(MASK) |
| rate = unknown.sum(-1).float() / (tokens.ne(PAD) & tokens.ne(BOS)).sum(-1).clamp_min(1) |
| with amp(device, precision): |
| logits = model(tokens, rate)[0].float() |
| logits[:, [PAD, BOS, MASK, UNK]] = -torch.inf |
| if not allow_eos: |
| logits[:, EOS] = -torch.inf |
| |
| |
| confidence_probs = logits.softmax(-1) |
| if temperature == 0: |
| candidates = logits.argmax(-1) |
| else: |
| values, indices = logits.topk(min(top_k, logits.shape[-1]), -1) |
| choice = torch.multinomial((values / temperature).softmax(-1), 1, generator=generator) |
| candidates = indices.gather(-1, choice).squeeze(-1) |
| confidence = confidence_probs.gather(-1, candidates[:, None]).squeeze(-1).masked_fill(~unknown[0], -1) |
| |
| remaining = int(unknown.sum()) |
| target = min(remaining - 1, math.floor(count * math.cos((step + 1) / iterations * math.pi / 2))) |
| if step == iterations - 1: |
| target = 0 |
| commit = remaining - max(0, target) |
| if policy == "adaptive": |
| commit = max(commit, int((confidence >= threshold).sum())) |
| ranking = confidence if policy != "left_to_right" else -torch.arange(len(ids), device=device).float().masked_fill(~unknown[0], float("inf")) |
| positions = ranking.topk(commit).indices |
| tokens[0, positions] = candidates[positions] |
| trace.append({"step": step + 1, "remaining": int(tokens.eq(MASK).sum()), |
| "text": tokenizer.decode(tokens[0].tolist(), skip_special_tokens=False)}) |
| if not tokens.eq(MASK).any(): |
| break |
| result = tokens[0].tolist() |
| return {"text": tokenizer.decode(result), "tokens": result, "trace": trace, |
| "forward_passes": len(trace), "latency_ms": (time.perf_counter() - started) * 1000, |
| "seed": seed} |
|
|
|
|
| def infill(model, tokenizer, prefix, suffix, missing_tokens=3, **kwargs): |
| if not 1 <= missing_tokens <= 64: |
| raise ValueError("missing_tokens must be 1..64") |
| ids = [BOS] + tokenizer.encode(prefix).ids + [MASK] * missing_tokens + tokenizer.encode(suffix).ids + [EOS] |
| return denoise(model, tokenizer, ids, **kwargs) |
|
|
|
|
| def continue_text(model, tokenizer, prompt, max_new_tokens=64, block_size=16, **kwargs): |
| if not 1 <= max_new_tokens <= 256 or not 1 <= block_size <= 64: |
| raise ValueError("max_new_tokens: 1..256; block_size: 1..64") |
| context = [BOS] + tokenizer.encode(prompt).ids |
| if len(context) + min(block_size, max_new_tokens) > model.config.max_length: |
| raise ValueError("Prompt too long for this model context") |
| output, trace, passes, latency = list(context), [], 0, 0. |
| base_seed = kwargs.pop("seed", 42) |
| for offset in range(0, max_new_tokens, block_size): |
| length = min(block_size, max_new_tokens - offset) |
| |
| window = [BOS] + output[1:][-(model.config.max_length - length - 1):] |
| result = denoise(model, tokenizer, window + [MASK] * length, seed=base_seed + offset, **kwargs) |
| new = result["tokens"][-length:] |
| passes += result["forward_passes"] |
| latency += result["latency_ms"] |
| trace.extend({**row, "block": offset // block_size + 1} for row in result["trace"]) |
| if EOS in new: |
| output.extend(new[:new.index(EOS)]) |
| break |
| output.extend(new) |
| return {"text": tokenizer.decode(output), "tokens": output, "trace": trace, "forward_passes": passes, |
| "latency_ms": latency, "seed": base_seed} |
|
|
