ilm.predictor_ilm

ILMPredictorUtilitiesMixin

clean_up_pred_ids(pred_ids, hold_mask=None)

Remove mask tokens inserted due to batched prediction.

ILMPredictor

Bases: Module, ILMPredictorUtilitiesMixin, Predictor[ILMBatch, ILMPredictionDict]

__init__(max_steps, max_length, tokenizer=None, noise_schedule=None, tokens_to_suppress=None, return_history=False, sampling_method='sample', top=1000, p=0.9, second_sampling_method=None, second_top=1000, second_p=0.9, model=None, input_constraint=False)

Constructor for ILMPredictor.

Parameters:

Name	Type	Description	Default
max_steps	int	The maximum number of steps to take.	required
max_length	int	The maximum length (excluding special tokens like PAD and MASK) of the generated text.	required
stopping_threshold	float	The threshold for stopping use on the length classification scores.	required
tokenizer	Tokenizer	The tokenizer. Typically, set after initialization but before calling predict.	None
noise_schedule	NoiseSchedule	The noise schedule. Typically, set after initialization but before calling predict.	None
tokens_to_suppress	List[str]	The tokens to suppress during generation.	None
return_history	bool	Whether to return the history.	False
sampling_method	Literal['sample', 'sample_top_k', 'sample_top_p']	The sampling method. When second_sampling_method is not provided, the specified method here is used to sample from the joint distribution of positions and tokens. When second_sampling_method is provided, the specified method here is used to sample from the token distribution (conditional) given the postions sampled using the second_sampling_method. "sample" means vanilla sampling from the distribution. "sample_top_k" means sampling from the top-k distribution. "sample_top_p" means sampling from the top-p distribution (neuclius samplingn).	'sample'
top	int	The top-k sampling parameter for sampling_method.	1000
p	float	The top-p sampling parameter for sampling_method.	0.9
second_sampling_method	Optional[Literal['sample', 'sample_top_k', 'sample_top_p']]	The second sampling method.	None
second_top	int	The second top-k sampling parameter for second_sampling_method.	1000
second_p	float	The second top-p sampling parameter for second_sampling_method.	0.9
model	Optional[ILMModel]	The model. Typically, set after initialization but before calling predict.	None

ILMPredictorWithLengthClassification

Bases: Module, ILMPredictorUtilitiesMixin, Predictor[ILMBatch, ILMPredictionDict]

__init__(max_steps, max_length, stopping_threshold=0.5, tokenizer=None, noise_schedule=None, tokens_to_suppress=None, return_history=False, sampling_method='sample', top=1000, p=0.9, second_sampling_method=None, second_top=1000, second_p=0.9, model=None, force_predict_first_step=False, input_constraint=False, use_high_precision=False, stopping_temperature=1.0)

Constructor for ILMPredictor.

Parameters:

Name	Type	Description	Default
max_steps	int	The maximum number of steps to take.	required
max_length	int	The maximum length (excluding special tokens like PAD and MASK) of the generated text.	required
stopping_threshold	float	The threshold for stopping use on the length classification scores.	0.5
tokenizer	Tokenizer	The tokenizer. Typically, set after initialization but before calling predict.	None
noise_schedule	NoiseSchedule	The noise schedule. Typically, set after initialization but before calling predict.	None
tokens_to_suppress	List[str]	The tokens to suppress during generation.	None
return_history	bool	Whether to return the history.	False
sampling_method	Literal['sample', 'sample_top_k', 'sample_top_p']	The sampling method. When second_sampling_method is not provided, the specified method here is used to sample from the joint distribution of positions and tokens. When second_sampling_method is provided, the specified method here is used to sample from the token distribution (conditional) given the postions sampled using the second_sampling_method. "sample" means vanilla sampling from the distribution. "sample_top_k" means sampling from the top-k distribution. "sample_top_p" means sampling from the top-p distribution (neuclius samplingn).	'sample'
top	int	The top-k sampling parameter for sampling_method.	1000
p	float	The top-p sampling parameter for sampling_method.	0.9
second_sampling_method	Optional[Literal['sample', 'sample_top_k', 'sample_top_p']]	The second sampling method.	None
second_top	int	The second top-k sampling parameter for second_sampling_method.	1000
second_p	float	The second top-p sampling parameter for second_sampling_method.	0.9
model	Optional[ILMModel]	The model. Typically, set after initialization but before calling predict.	None