Minimal code to train ELMo models in TensorFlow.
Heavily based on https://github.com/allenai/bilm-tf .
Most changes are simplifications and updating the code to the recent versions of TensorFlow 1.
NB: TensorFlow is outdated and has severe vulnerabilities. The code in this repository, however, is not updated to use TensorFlow 2. Thus, you take the risk.
See also our repository with simple code to infer contextualized word vectors from pre-trained ELMo models.
python3 bilm/train_elmo.py --train_prefix $DATA --size $SIZE --vocab_file $VOCAB --save_dir $OUT
where
$DATA is a path to the directory containing 2 or more of (possibly gzipped) plain text files: your training corpus.
$SIZE if the number of word tokens in $DATA (necessary to properly construct and log batches).
$VOCAB is a (possibly gzipped) one-word-per-line vocabulary file to be used for language modeling; it should always contain at least <S>, </S> and <UNK>.
$OUT is a directory where the TensorFlow checkpoints will be saved.
Before training, please review the settings in bilm/train_elmo.py. The most important are:
- batch_size (default 128)
- n_gpus (default 2; if no GPU, all available CPU cores are used)
- LSTM dimensionality (default 2048; the original paper used 4096)
- n_epochs (default 3; optimal value depends on the size of your corpus)
- n_negative_samples_batch (default 4096; the original paper used 8192)
After the training, use the bilm/dump_weights.py script to convert the checkpoints to and HDF5 model.
python3 bilm/dump_weights.py --save_dir $MODEL_DIR --outfile $MODEL_DIR/model.hdf5
Save your vocabulary file in the same directory.
Change the n_characters value in the options.json file from 261 to 262 to use the saved model for inference.
More details at https://github.com/allenai/bilm-tf