generate.py — Checkpoint text and token generation

scripts/generate.py exposes generation from compatible Kimi K3 trainer checkpoints without imposing the original evaluation protocol. It is an interface around generate_tokens(), not a separate decoding engine. The implementation details are documented in inference.py. Two input paths matter: text prompts use a Hugging Face tokenizer; raw token IDs avoid loading a tokenizer, making an offline checkpoint smoke test possible.

Code excerpts below are dedented from their enclosing functions for readability; the displayed statements and ordering are preserved.

CLI selection and defaults

The parser requires --checkpoint and exactly one of --prompt TEXT and --token-ids COMMA_SEPARATED_IDS (an argparse mutually exclusive group).

For an illustrative offline token-prompt run, replace the path with a locally available compatible checkpoint:

python scripts/generate.py --checkpoint path/to/last.pt --token-ids 7,13,29 --device cpu --mode greedy

--token-ids means no tokenizer is loaded, not that the model ignores the vocabulary and context constraints. Conversely, --prompt "..." uses the named tokenizer at the specified revision and may need cached files or network access unless --local-files-only is set.

Convert the requested prompt to IDs

The first branch of main() resolves the device and rejects BF16 on non-CUDA devices. With raw input, it splits on commas, strips whitespace, converts each field to int, and wraps a conversion ValueError with a flag-specific message:

if args.token_ids is not None:
    try:
        ids = [int(x.strip()) for x in args.token_ids.split(",")]
    except ValueError as exc:
        raise ValueError("--token-ids must be comma-separated integers") from exc

Syntax parsing does not itself range-check the token IDs. That happens later in the inference service, where empty prompts, negative/out-of-vocabulary IDs and overlong contexts are rejected. In text mode, AutoTokenizer.from_pretrained() uses use_fast=True and tokenizer.encode(..., add_special_tokens=False). In particular, the CLI does not prepend a BOS token on its own.

Load, validate, then generate

The checkpoint and EOS branch is intentionally after prompt encoding:

model, metadata = load_checkpoint_model(args.checkpoint, device=device, backend=args.backend)
if tokenizer is not None and len(tokenizer) != model.config.vocab_size:
    raise ValueError("Tokenizer vocabulary does not match checkpoint")
eos_id = args.eos_id if args.eos_id is not None else (
    tokenizer.eos_token_id if tokenizer is not None and tokenizer.eos_token_id is not None else 0
)

In text mode, it checks tokenizer length against checkpoint vocabulary size. In raw-ID mode no tokenizer-vocabulary equality can be checked. By default, EOS comes from the tokenizer if available, otherwise ID zero. --eos-id takes precedence, including if set to zero; the inference service verifies that the chosen ID is in range. The checkpoint loader requires format-version-7 model/config state, checks strict weight assignment, restores embedding/LM-head tying and sets evaluation mode. The fused backend requires CUDA; BF16 also needs supported CUDA hardware, not just a CUDA device name.

The CLI passes mode, temperature, top-k, top-p, seed and maximum length to generate_tokens(). That routine recomputes next-token logits on the accumulated prefix on every step, pads a reference-KDA prefix to a full chunk if necessary, applies argmax or CPU-seeded sampling, and stops early at EOS. It has no sliding context window or KV-cache reuse. Even with early stopping, it rejects requests where len(prompt_ids) + max_new_tokens exceeds model max_seq_len (and chunk padding may impose another limit). Greedy mode ignores the sampling parameters; in sample mode invalid temperature/top-k/top-p raise ValueError.

JSON contract and side effects

After generation, the final block prints an indented JSON object containing checkpoint metadata, selected device, precision, prompt_token_ids and generated_token_ids. The latter contains new tokens only, not the entire prompt-plus-completion sequence. Text input adds generated_text, decoded only from the new IDs, with special tokens skipped.

if tokenizer is not None:
    result["generated_text"] = tokenizer.decode(generated, skip_special_tokens=True)
print(json.dumps(result, indent=2, ensure_ascii=False))

There is no --output option, persistent generation state, checkpoint modification or dataset access. Loading local weights is read-only; tokenizer retrieval may populate an external Hugging Face cache. Model device memory usage can be substantial, and sampling is seeded only at the CPU token-selection generator level; the CLI does not claim full-system deterministic execution.

Verification and limits

The inspected test_public_inference.py constructs a tiny trainer-format checkpoint, checks weight tying and finite next-token logits, asserts that a one-token greedy continuation matches argmax, rejects an overlong request, and tests selection with a fixed top-k and invalid top-p. Those assertions verify core primitives, not the text-tokenizer network path or a subprocess run of this CLI. All descriptions here are based on pinned source and inspected tests, not a new generation trial.