kiwi.lib.search
RangeConfig
Specify a continuous interval, or a discrete range when step is set.
ClassWeightsConfig
Specify the range to search in for the tag loss weights.
SearchOptions
Base class for all pydantic configs. Used to configure base behaviour of configs.
Configuration
Objective
The objective to be optimized by the Optuna hyperparameter search.
search_from_file(filename: Path)
search_from_file
Load options from a config file and calls the training procedure.
search_from_configuration(configuration_dict: dict)
search_from_configuration
Run the entire training pipeline using the configuration options received.
get_suggestion(trial, param_name: str, config: Union[List, RangeConfig]) → Union[bool, float, int]
get_suggestion
Let the Optuna trial suggest a parameter value with name param_name
param_name
setup_run(directory: Path, seed: int, debug=False, quiet=False) → Path
setup_run
Set up the output directory structure for the Optuna search outputs.
run(config: Configuration)
run
Run hyperparameter search according to the search configuration.
kiwi.lib.search.
logger
Bases: kiwi.utils.io.BaseConfig
kiwi.utils.io.BaseConfig
lower
The lower bound of the search range.
upper
The upper bound of the search range.
step
Specify a step size to create a discrete range of search values.
distribution
Specify the distribution over the search range. Uniform is recommended for all hyperparameters except for the learning rate, for which loguniform is recommended. Only works when a RangeConfig without step is specified.
target_tags
Loss weight for the target tags.
gap_tags
Loss weight for the gap tags.
source_tags
Loss weight for the source tags.
patience
Number of training validations without improvement to wait before stopping training.
validation_steps
Rely on the Kiwi training options to early stop bad models.
search_mlp
To use or not to use an MLP after the encoder.
search_word_level
Try with and without word level output. Useful to figure out if word level prediction is helping HTER regression performance.
search_hter
Try with and without sentence level output. Useful to figure out if HTER regression is helping word level performance.
learning_rate
Search the learning rate value.
dropout
Search the dropout rate used in the decoder.
warmup_steps
Search the number of steps to warm up the learning rate.
freeze_epochs
Search the number of epochs to freeze the encoder.
class_weights
Search the word-level tag loss weights.
sentence_loss_weight
Search the weight to scale the sentence loss objective with.
hidden_size
Search the hidden size of the MLP decoder.
bottleneck_size
Search the size of the hidden layer in the decoder bottleneck.
search_method
Use random search or the (multivariate) Tree-structured Parzen Estimator, or shorthand: TPE. See optuna.samplers for more details about these methods.
optuna.samplers
check_consistency
base_config
Kiwi train configuration used as a base to configure the search models. Can be a path or a yaml configuration properly indented under this argument.
directory
Output directory.
seed
Make the search reproducible.
search_name
The name used by the Optuna MLflow integration. If None, Optuna will create a unique hashed name.
num_trials
The number of search trials to run.
num_models_to_keep
The number of model checkpoints that are kept after finishing search. The best checkpoints are kept, the others removed to free up space. Keep all model checkpoints by setting this to -1.
options
Configure the search method and parameter ranges.
load_study
Continue from a previous saved study, i.e. from a study.pkl file.
study.pkl
verbose
quiet
parse_base_config
filename – of the configuration file.
an object with training information.
configuration_dict – dictionary with options.
object with training information.
Let the Optuna trial suggest a parameter value with name param_name based on the range configuration.
trial – an Optuna trial
param_name (str) – the name of the parameter to suggest a value for
config (Union[List, RangeConfig]) – the parameter search space
The suggested parameter value.
The call method initializes a Kiwi training config based on Optuna parameter suggestions, trains Kiwi, and then returns the output.
The model paths of the models are saved internally together with the objective value obtained for that model. These can be used to prune model checkpoints after completion of the search.
config (Configuration) – the search configuration.
base_config_dict (dict) – the training configuration to serve as base, in dictionary form.
main_metric
The main validation metric as it is formatted by the Kiwi trainer.
This can be used to access the main metric value after training via train_info.best_metrics[objective.main_metric].
train_info.best_metrics[objective.main_metric]
num_train_lines
The number of lines in the training data.
updates_per_epochs
The number of parameter updates per epochs.
best_model_paths
Return the model paths sorted from high to low by their objective score.
best_train_configs
Return the train configs sorted from high to low by their objective score.
prune_models
Keep only the best model checkpoints and remove the rest to free up space.
suggest_train_config
Use the trial to suggest values to initialize a training configuration.
trial – An Optuna trial to make hyperparameter suggestions.
A Kiwi train configuration and a dictionary with the suggested Optuna parameter names and values that were set in the train config.
__call__
Train Kiwi with the hyperparameter values suggested by the trial and return the value of the main metric.
A float with the value obtained by the Kiwi model, as measured by the main metric configured for the model.
config (Configuration) – search configuration
an optuna study summarizing the search results
kiwi.lib.pretrain
kiwi.lib.train