File size: 5,312 Bytes
c6f3061 | 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 | """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
# Ranking always uses the untempered distribution. Using T≈0 here
# would turn every argmax confidence into 1 and destroy its ordering.
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)
# Cosine schedule keeps more holes early and commits all by the last pass.
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)
# Keep BOS as a boundary when a continuation exceeds the context window.
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}
|