1b53c49e1a
* diffusion training plugin * cleanup * nits * fixes + improvements * add back in reinit_weights (clobbered?); masking / pretrain fixes * nits * cleanup; tests draft * sample generation, tests fixes * fixes * nits * add inference support; add auto-mask token support * nits * nits * progress * simplify logging * lint * prefix args with diffusion_ * coderabbito * tests fix * nit * nits * cleanup + nits * nits * fix SFT sample gen * fixes * fix * comments * comments * lint * reward model lora fix * cleanup; fix pretraining_dataset case * gradio inference * update cfgs * update cfgs * train, generation parity, cleanup * fix * simplify * test * test fix
20 lines
556 B
Python
20 lines
556 B
Python
"""Diffusion LM training plugin init."""
|
|
|
|
from .args import DiffusionArgs, DiffusionConfig
|
|
from .callbacks import DiffusionGenerationCallback
|
|
from .generation import generate
|
|
from .plugin import DiffusionPlugin
|
|
from .trainer import DiffusionTrainer
|
|
from .utils import create_bidirectional_attention_mask, resolve_mask_token_id
|
|
|
|
__all__ = [
|
|
"DiffusionArgs",
|
|
"DiffusionPlugin",
|
|
"DiffusionTrainer",
|
|
"generate",
|
|
"resolve_mask_token_id",
|
|
"create_bidirectional_attention_mask",
|
|
"DiffusionGenerationCallback",
|
|
"DiffusionConfig",
|
|
]
|