Skip to main content

Utilities to use allennlp with wandb

Project description



Utilities and boilerplate code which allows using Weights & Biases to tune the hypereparameters for any AllenNLP model without a single line of extra code!

What does it do?

  1. Log a single run or a hyperparameter search sweep without any extra code, just using configuration files.

  2. Use Weights & Biases' bayesian hyperparameter search engine + hyperband in any AllenNLP project.

Quick start


$ pip install wandb-allennlp
$ echo wandb_allennlp >> .allennlp_plugins

Log a single run

  1. Create your model using AllenNLP along with a training configuration file as you would normally do.

  2. Add a trainer callback in your config file. Use one of the following based on your AllenNLP version:


trainer: {
    type: 'callback',
    callbacks: [
        type: 'wandb_allennlp',
        files_to_save: ['config.json'],
        files_to_save_at_end: ['*.tar.gz'],
  1. Execute the allennlp train-with-wandb command instead of allennlp train. It supports all the arguments present in allennlp train. However, the --overrides have to be specified in the --kw value or --kw=value form, where kw is the parameter to override and value is its value. Use the dot notation for nested parameters. For instance, {'model': {'embedder': {'type': xyz}}} can be provided as --model.embedder.type xyz.
allennlp  train-with-wandb model_configs/my_config.jsonnet --include-package=package_with_my_registered_classes --include-package=another_package --wandb-run-name=my_first_run --wandb-tags=any,set,of,non-unique,tags,that,identify,the,run,without,spaces

Hyperparameter Search

  1. Create your model using AllenNLP along with a training configuration file as you would normally do. For example:
local data_path = std.extVar('DATA_PATH');
local a = std.parseJson(std.extVar('a'));
local bool_value = std.parseJson(std.extVar('bool_value'));
local int_value = std.parseJson(std.extVar('int_value'));

  type: 'train_test_log_to_wandb',
  evaluate_on_test: true,
  dataset_reader: {
    type: 'snli',
    token_indexers: {
      tokens: {
        type: 'single_id',
        lowercase_tokens: true,
  train_data_path: data_path + '/snli_1.0_test/snli_1.0_train.jsonl',
  validation_data_path: data_path + '/snli_1.0_test/snli_1.0_dev.jsonl',
  test_data_path: data_path + '/snli_1.0_test/snli_1.0_test.jsonl',
  model: {
    type: 'parameter-tying',
    a: a,
    b: a,
    d: 0,
    bool_value: bool_value,
    bool_value_not: !bool_value,
    int_value: int_value,
    int_value_10: int_value + 10,

  data_loader: {
    batch_sampler: {
      type: 'bucket',
      batch_size: 64,
  trainer: {
    optimizer: {
      type: 'adam',
      lr: 0.001,
      weight_decay: 0.0,
    cuda_device: -1,
    num_epochs: 2,
    callbacks: [
        type: 'wandb_allennlp',
        files_to_save: ['config.json'],
        files_to_save_at_end: ['*.tar.gz'],
  1. Create a sweep configuration file and generate a sweep on the wandb server. Note that the tied parameters that are accepted through environment variables are specified using the prefix env. in the sweep config. For example:
name: parameter_tying_test_console_script_v0.2.4
program: allennlp
  - ${program} #omit the interpreter as we use allennlp train command directly
  - "train-with-wandb" # subcommand
  - "configs/parameter_tying_v0.2.4.jsonnet"
  - "--include-package=models" # add all packages containing your registered classes here
  - "--include-package=allennlp_models"
  - ${args}
method: bayes
  name: training_loss
  goal: minimize
  # hyperparameters start with overrides
  # Ranges
  # Add env. to tell that it is a top level parameter
    min: 1
    max: 10
    distribution: uniform
    values: [true, false]
    values: [-1, 0, 1, 10]
    value: 1
  1. Create the sweep on wandb.
$ wandb sweep path_to_sweep.yaml
  1. Set the other environment variables required by your jsonnet.
export DATA_DIR=./data
  1. Start the search agents.
wandb agent <sweep_id>

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

wandb_allennlp-0.3.3.tar.gz (16.9 kB view hashes)

Uploaded source

Built Distribution

wandb_allennlp-0.3.3-py3-none-any.whl (19.1 kB view hashes)

Uploaded py3

Supported by

AWS AWS Cloud computing Datadog Datadog Monitoring Fastly Fastly CDN Google Google Object Storage and Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page