| |
| """Text continuation for KerzgrAI/Rondo-Exp and KerzgrAI/Monody-Exp. |
| |
| Both are proof-of-concept base models: they continue text and follow no |
| instructions. Rondo-Exp is masked diffusion: the continuation starts as masks |
| and is revealed in random order, with up to --drafts tokens drafted and checked |
| per forward (exact self-speculative decoding). Monody-Exp is its |
| autoregressive twin and samples one token per forward with a cache. |
| |
| Run from a downloaded repo, which this script reads its model from: |
| python inference.py --prompt "The water cycle is" |
| python inference.py --prompt "Photosynthesis" --max-new-tokens 128 --temperature 0.8 |
| Or anywhere, downloading the model first: |
| python inference.py --repo KerzgrAI/Monody-Exp --prompt "The water cycle is" |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib |
| import subprocess |
| import sys |
| import time |
| import warnings |
| from pathlib import Path |
| from types import ModuleType |
|
|
| import torch |
| from huggingface_hub import snapshot_download |
| from safetensors.torch import load_file |
| from tokenizers import Tokenizer |
|
|
| REPOS = ("KerzgrAI/Rondo-Exp", "KerzgrAI/Monody-Exp") |
| |
| FLA_COMMIT = "cbb0a72efb55c18ca0ef4f298298317573ad2cb3" |
| FLA_REPO = "https://github.com/fla-org/flash-linear-attention.git" |
| FLA_MODULES = ( |
| "fla.ops.gdn2", |
| "fla.layers.gdn2", |
| "fla.modules.rotary", |
| "fla.ops.parallax.parallel", |
| ) |
| TINY_GDN_FILES = ( |
| "tiny_gdn/__init__.py", |
| "tiny_gdn/config.py", |
| "tiny_gdn/detokenize.py", |
| "tiny_gdn/fugue.py", |
| "tiny_gdn/fugue_sampling.py", |
| "tiny_gdn/fused_ops.py", |
| "tiny_gdn/gated_residual_ops.py", |
| "tiny_gdn/haiku_layers.py", |
| "tiny_gdn/kdn/__init__.py", |
| "tiny_gdn/kdn/kernels.py", |
| "tiny_gdn/kdn/layer.py", |
| "tiny_gdn/kdn/recurrence.py", |
| "tiny_gdn/model.py", |
| "tiny_gdn/ngram.py", |
| "tiny_gdn/nn_common.py", |
| "tiny_gdn/parallax_ops.py", |
| "tiny_gdn/senryu2.py", |
| ) |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser( |
| description="Rondo-Exp / Monody-Exp text continuation" |
| ) |
| parser.add_argument("--prompt", required=True, help="Text to continue") |
| parser.add_argument( |
| "--repo", |
| default=None, |
| help=f"Hub repo or local folder to load ({' or '.join(REPOS)}); default: this script's folder", |
| ) |
| parser.add_argument("--max-new-tokens", type=int, default=256) |
| parser.add_argument("--temperature", type=float, default=1.0, help="0 is greedy") |
| parser.add_argument("--top-p", type=float, default=0.9) |
| parser.add_argument("--top-k", type=int, default=0, help="0 keeps every token") |
| parser.add_argument("--seed", type=int, default=0) |
| parser.add_argument( |
| "--drafts", |
| type=int, |
| default=8, |
| help="Rondo-Exp: tokens drafted per forward; 0 reveals one token per forward", |
| ) |
| args = parser.parse_args() |
| if args.max_new_tokens <= 0: |
| parser.error("--max-new-tokens must be positive") |
| if ( |
| args.temperature < 0 |
| or not 0 < args.top_p <= 1 |
| or args.top_k < 0 |
| or args.drafts < 0 |
| ): |
| parser.error( |
| "--temperature >= 0, 0 < --top-p <= 1, --top-k >= 0 and --drafts >= 0" |
| ) |
| return args |
|
|
|
|
| def model_directory(repo: str | None) -> Path: |
| here = Path(__file__).resolve().parent |
| if repo is None: |
| if (here / "config.json").is_file(): |
| return here |
| raise SystemExit( |
| f"No model next to {Path(__file__).name}; pass --repo {REPOS[0]} or --repo {REPOS[1]}" |
| ) |
| if Path(repo).is_dir(): |
| return Path(repo).resolve() |
| return Path(snapshot_download(repo)) |
|
|
|
|
| def ensure_fla() -> None: |
| try: |
| for name in FLA_MODULES: |
| importlib.import_module(name) |
| return |
| except Exception as error: |
| print( |
| f"Installing flash-linear-attention {FLA_COMMIT[:7]} ({error}) …", |
| file=sys.stderr, |
| flush=True, |
| ) |
| pip = [sys.executable, "-m", "pip", "install"] |
| subprocess.check_call([*pip, "einops", "transformers>=4.45.0"]) |
| subprocess.check_call( |
| [*pip, "--no-deps", "--force-reinstall", f"git+{FLA_REPO}@{FLA_COMMIT}"] |
| ) |
| for name in [ |
| name for name in sys.modules if name == "fla" or name.startswith("fla.") |
| ]: |
| del sys.modules[name] |
| importlib.invalidate_caches() |
| for name in FLA_MODULES: |
| importlib.import_module(name) |
|
|
|
|
| def load_model( |
| directory: Path, device: torch.device |
| ) -> tuple[torch.nn.Module, ModuleType | None]: |
| """The model in ``directory`` on ``device`` in bfloat16, and fugue_sampling when it is Rondo.""" |
|
|
| missing = [name for name in TINY_GDN_FILES if not (directory / name).is_file()] |
| if missing: |
| raise SystemExit(f"{directory} lacks the bundled code: {', '.join(missing)}") |
| sys.path.insert(0, str(directory)) |
| config = importlib.import_module("tiny_gdn.config").TinyGDNConfig.from_json( |
| directory / "config.json" |
| ) |
| if config.model_type == "fugue": |
| model_class = importlib.import_module("tiny_gdn.fugue").FugueForMaskedDiffusion |
| sampling = importlib.import_module("tiny_gdn.fugue_sampling") |
| else: |
| model_class = importlib.import_module("tiny_gdn.senryu2").Senryu2ForCausalLM |
| sampling = None |
| with torch.device(device): |
| model = model_class(config) |
| model = model.to(dtype=torch.bfloat16) |
| model.load_state_dict( |
| load_file(directory / "model.safetensors", device="cpu"), strict=True |
| ) |
| return model.eval(), sampling |
|
|
|
|
| def sample( |
| logits: torch.Tensor, args: argparse.Namespace, generator: torch.Generator |
| ) -> int: |
| logits = logits.float() |
| if args.temperature == 0: |
| return int(logits.argmax()) |
| logits = logits / args.temperature |
| if 0 < args.top_k < logits.numel(): |
| logits = logits.masked_fill( |
| logits < torch.topk(logits, args.top_k).values[-1], -torch.inf |
| ) |
| if args.top_p < 1: |
| ordered, order = torch.sort(logits, descending=True) |
| drop = torch.softmax(ordered, dim=-1).cumsum(-1) > args.top_p |
| drop[1:] = drop[:-1].clone() |
| drop[0] = False |
| logits = logits.scatter(0, order, ordered.masked_fill(drop, -torch.inf)) |
| return int(torch.multinomial(torch.softmax(logits, dim=-1), 1, generator=generator)) |
|
|
|
|
| @torch.inference_mode() |
| def continue_autoregressive(model, prompt: list[int], args, generator, emit) -> int: |
| eos = int(model.config.eos_token_id) |
| room = model.config.max_position_embeddings - len(prompt) |
| output = model( |
| torch.tensor([prompt], device=generator.device), |
| return_logits=True, |
| logits_to_keep=1, |
| use_cache=True, |
| ) |
| produced = 0 |
| for _ in range(min(args.max_new_tokens, room)): |
| token = sample(output.logits[0, -1], args, generator) |
| if token == eos: |
| break |
| emit(token) |
| produced += 1 |
| output = model( |
| torch.tensor([[token]], device=generator.device), |
| return_logits=True, |
| logits_to_keep=1, |
| past_key_values=output.past_key_values, |
| use_cache=True, |
| ) |
| return produced |
|
|
|
|
| @torch.inference_mode() |
| def continue_diffusion( |
| model, sampling: ModuleType, prompt: list[int], args, generator, emit |
| ) -> int: |
| eos = int(model.config.eos_token_id) |
| mask = model.mask_token_id |
| settings = sampling.DiffusionSampling( |
| new_tokens=args.max_new_tokens, |
| block_length=None, |
| threshold=None, |
| order="random", |
| steps_per_block=None if args.drafts else args.max_new_tokens, |
| draft_tokens=args.drafts, |
| temperature=args.temperature, |
| top_k=args.top_k, |
| top_p=args.top_p, |
| ) |
| shown = 0 |
| forwards = 0 |
|
|
| def show(canvas: torch.Tensor) -> bool: |
| nonlocal shown |
| committed = sampling.committed_prefix(canvas, mask).tolist() |
| for token in committed[shown:]: |
| if token == eos: |
| return True |
| emit(token) |
| shown += 1 |
| return False |
|
|
| def on_step(state) -> bool: |
| nonlocal forwards |
| forwards = state.step |
| return show(state.tokens[0]) |
|
|
| canvas = sampling.diffusion_generate( |
| model, |
| torch.tensor([prompt], device=generator.device), |
| settings, |
| generator=generator, |
| eos_token_id=eos, |
| on_step=on_step, |
| ) |
| show(canvas[0]) |
| print(f"\n[{forwards} forwards]", file=sys.stderr, end="", flush=True) |
| return shown |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| warnings.filterwarnings("ignore", message=r"tl\.make_block_ptr is deprecated") |
| if not torch.cuda.is_available(): |
| raise SystemExit( |
| "An NVIDIA GPU is required: the GDN-2 kernels run in Triton on CUDA." |
| ) |
| device = torch.device("cuda") |
| directory = model_directory(args.repo) |
| ensure_fla() |
| model, sampling = load_model(directory, device) |
| tokenizer = Tokenizer.from_file(str(directory / "tokenizer.json")) |
| decoder = importlib.import_module("tiny_gdn.detokenize").IncrementalUtf8Decoder( |
| tokenizer |
| ) |
| prompt = tokenizer.encode(args.prompt, add_special_tokens=False).ids |
| if not prompt: |
| raise SystemExit("The prompt produced no tokens") |
| room = model.config.max_position_embeddings - args.max_new_tokens |
| if room < 1: |
| raise SystemExit( |
| f"--max-new-tokens leaves no room in {model.config.max_position_embeddings} positions" |
| ) |
| prompt = prompt[-room:] |
| generator = torch.Generator(device=device).manual_seed(args.seed) |
|
|
| def emit(token: int) -> None: |
| print(decoder.push(token), end="", flush=True) |
|
|
| print(args.prompt, end="", flush=True) |
| started = time.perf_counter() |
| if sampling is None: |
| produced = continue_autoregressive(model, prompt, args, generator, emit) |
| else: |
| produced = continue_diffusion(model, sampling, prompt, args, generator, emit) |
| shown = len(decoder.text) |
| print(decoder.finalize()[shown:], end="") |
| elapsed = time.perf_counter() - started |
| print(f"\n[{produced} tokens in {elapsed:.1f}s]", file=sys.stderr, flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
