Function main

examples/discrete_diffusion/sample_llada2.py:36–259

()

def main():
    parser = argparse.ArgumentParser(
        description="Generate text using LLaDA2Pipeline with block-wise discrete diffusion."
    )
    parser.add_argument(
        "--model_id",
        type=str,
        default="inclusionAI/LLaDA2.0-mini",
        help="HuggingFace model ID or path to local model.",
    )
    parser.add_argument(
        "--prompt",
        type=str,
        default="Why does Camus think that Sisyphus is happy?",
        help="Text prompt to generate from.",
    )
    parser.add_argument(
        "--gen_length",
        type=int,
        default=2048,
        help="Number of tokens to generate.",
    )
    parser.add_argument(
        "--block_length",
        type=int,
        default=32,
        help="Size of each generation block.",
    )
    parser.add_argument(
        "--num_inference_steps",
        type=int,
        default=32,
        help="Number of refinement steps per block.",
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.0,
        help="Sampling temperature (0.0 for greedy).",
    )
    parser.add_argument(
        "--top_p",
        type=float,
        default=None,
        help="Nucleus sampling probability threshold.",
    )
    parser.add_argument(
        "--top_k",
        type=int,
        default=None,
        help="Top-k sampling parameter.",
    )
    parser.add_argument(
        "--threshold",
        type=float,
        default=0.95,
        help="Confidence threshold for committing tokens.",
    )