| |
| """Sample a pinned Di4C language student from local server assets only.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| from _distilled_runtime import ( |
| benchmark_model, |
| checkpoint_state, |
| configure_for_sampling, |
| install_upstream, |
| load_config, |
| load_tokenizer, |
| load_tokenizer_binding, |
| materialize_model, |
| offline_huggingface, |
| parse_bool, |
| require_directory, |
| require_file, |
| seed_everything, |
| validate_embedded_config, |
| validate_sampling_config, |
| write_capture_atomic, |
| ) |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--upstream-root", type=Path, required=True) |
| parser.add_argument("--checkpoint", type=Path, required=True) |
| parser.add_argument("--config", type=Path, required=True) |
| parser.add_argument("--checkpoint-sha256", required=True) |
| parser.add_argument("--config-sha256", required=True) |
| parser.add_argument("--data-config", type=Path, required=True) |
| parser.add_argument("--downloads-manifest", type=Path, required=True) |
| parser.add_argument("--tokenizer-snapshot", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--sample-count", type=int, required=True) |
| parser.add_argument("--batch-size", type=int, required=True) |
| parser.add_argument("--num-steps", type=int, required=True) |
| parser.add_argument("--seq-len", type=int, required=True) |
| parser.add_argument("--seed", type=int, required=True) |
| parser.add_argument("--sampler", choices=("ancestral",), required=True) |
| parser.add_argument( |
| "--teacher-family", choices=("masked_mdlm", "uniform_duo"), required=True |
| ) |
| parser.add_argument("--dataset", choices=("lm1b", "owt"), required=True) |
| parser.add_argument("--generation-mode", choices=("unconditional", "conditional_prefix"), default="unconditional") |
| parser.add_argument("--conditioning-manifest", type=Path) |
| parser.add_argument("--conditioning-manifest-sha256") |
| parser.add_argument("--conditioning-config-sha256") |
| parser.add_argument("--prefix-length", type=int) |
| parser.add_argument("--evaluation-continuation-length", type=int) |
| parser.add_argument("--prompt-count", type=int) |
| parser.add_argument("--diversity-prompt-count", type=int) |
| parser.add_argument("--completions-per-diversity-prompt", type=int) |
| parser.add_argument("--completion-schedule") |
| parser.add_argument("--offline", type=parse_bool, required=True) |
| parser.add_argument( |
| "--allow-missing-embedded-config", type=parse_bool, default=False |
| ) |
| parser.add_argument("--benchmark-output", type=Path) |
| parser.add_argument("--benchmark-metadata", type=Path) |
| parser.add_argument("--benchmark-precision", choices=("author",)) |
| args = parser.parse_args(argv) |
| benchmark_values = (args.benchmark_output, args.benchmark_metadata, args.benchmark_precision) |
| if any(value is not None for value in benchmark_values) and not all( |
| value is not None for value in benchmark_values |
| ): |
| parser.error("benchmark arguments must be provided together") |
|
|
| upstream = require_directory(args.upstream_root, "pinned Di4C language source") |
| require_file(upstream / "src/sdtt/main.py", "pinned Di4C entrypoint") |
| install_upstream(upstream) |
| checkpoint = require_file(args.checkpoint, "Di4C student checkpoint") |
| config_path = require_file(args.config, "Di4C sampling config") |
| binding = load_tokenizer_binding( |
| args.data_config, |
| args.downloads_manifest, |
| args.dataset, |
| args.tokenizer_snapshot, |
| ) |
| with offline_huggingface(args.offline): |
| from sdtt.core.distill.multi_round_sdtt import MultiRoundSDTT |
|
|
| config = load_config(config_path, args.config_sha256) |
| state, embedded_config = checkpoint_state( |
| checkpoint, args.checkpoint_sha256 |
| ) |
| validate_sampling_config( |
| config, |
| binding=binding, |
| sequence_length=args.seq_len, |
| require_di4c=True, |
| ) |
| validate_embedded_config( |
| config, |
| embedded_config, |
| allow_missing_embedded_config=args.allow_missing_embedded_config, |
| ) |
| configure_for_sampling( |
| config, tokenizer_snapshot=binding.snapshot, seq_len=args.seq_len |
| ) |
| config.parameterization.checkpoint_path = str(checkpoint) |
| tokenizer = load_tokenizer(binding.snapshot) |
| seed_everything(args.seed) |
| model = materialize_model( |
| model_type=MultiRoundSDTT, |
| config=config, |
| tokenizer=tokenizer, |
| state=state, |
| strict=False, |
| ) |
| if args.benchmark_output is not None: |
| benchmark_model( |
| model=model, |
| output=args.benchmark_output, |
| metadata_path=args.benchmark_metadata, |
| precision=args.benchmark_precision, |
| num_steps=args.num_steps, |
| seq_len=args.seq_len, |
| sampler=args.sampler, |
| generation_mode=args.generation_mode, |
| conditioning_manifest=args.conditioning_manifest, |
| conditioning_manifest_sha256=args.conditioning_manifest_sha256, |
| prompt_count=args.prompt_count, |
| diversity_prompt_count=args.diversity_prompt_count, |
| completions_per_diversity_prompt=args.completions_per_diversity_prompt, |
| ) |
| else: |
| write_capture_atomic( |
| args.output, |
| model=model, |
| tokenizer=tokenizer, |
| sample_count=args.sample_count, |
| batch_size=args.batch_size, |
| num_steps=args.num_steps, |
| seq_len=args.seq_len, |
| sampler=args.sampler, |
| generation_mode=args.generation_mode, |
| conditioning_manifest=args.conditioning_manifest, |
| conditioning_manifest_sha256=args.conditioning_manifest_sha256, |
| prompt_count=args.prompt_count, |
| diversity_prompt_count=args.diversity_prompt_count, |
| completions_per_diversity_prompt=args.completions_per_diversity_prompt, |
| ) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
