mlm.predictor_mlm

MLMPredictor

Bases: Module, Predictor[MLMBatch, MLMPredictionDict]

Base predictor for MLM. Stochastically selects positions to unmask based on max_steps and max_new_tokens.

__init__(max_steps, max_new_tokens=None, tokenizer=None, model=None, noise_schedule=None, top_k=None, top_p=None, skip_special_tokens=True)

Initialize MLM Predictor.

Parameters:

Name	Type	Description	Default
max_steps	int	Maximum number of prediction steps.	required
tokenizer	Optional[Tokenizer]	Tokenizer for encoding/decoding.	None
noise_schedule	Optional[NoiseSchedule]	Noise schedule for the diffusion process.	None
top_k	Optional[int]	Top-k sampling parameter.	None
top_p	Optional[float]	Top-p sampling parameter.	None
model	Optional[MLMModel]	The MLM model to use for predictions.	None

decode(results)

Parameters:

Name	Type	Description	Default
results	MLMStepResults	x: Integer[TT, " batch seq_len"] Current predicted sequence.	required

Returns: out: List[str] Decoded sequence with special tokens. x: Integer[TT, " batch seq_len"] Current predicted sequence.