LPU-NN

Neural language processing models on PyTorch, built on LPU.

日本語版のドキュメントは README.ja.md にあります。

Status

This package revives a private research codebase written in 2019-2020 for reproducing and prototyping neural language processing models. It is being ported and modernized incrementally, so the API is not yet stable.

What currently runs end to end on PyTorch 2.x / Python 3.13:

• sequence-to-sequence training (lpu-nn-train-seq2seq)
• decoding with beam search (lpu-nn-run-seq2seq)
• sequence matching and ranking (lpu-nn-train-match-ranker, lpu-nn-run-match-ranker), with the RE2 and Compare-Aggregate poolers
• BERT pre-training (lpu-nn-train-bert) and fine-tuning for classification (lpu-nn-train-bert-classifier) and pair ranking (lpu-nn-train-bert-ranker)
• sequence tagging (lpu-nn-train-tagger), over a BiLSTM, a transformer or a BERT encoder, with a linear or a CRF decoder
• character language modeling (lpu-nn-train-embedding), the tokenizer command (lpu-nn-run-tokenizer) and an HTTP server for a trained sequence-to-sequence model (lpu-nn-serve-seq2seq)

The whole of the original codebase is ported.

Requirements

• Python 3.10 or later
• PyTorch 2.4 or later (a CUDA build is recommended for training)

Installation

$ pip install 'lpu-nn @ git+https://github.com/akivajp/lpu-nn.git'

For development:

$ uv sync

Training curves are written only when the optional plot extra is installed:

$ pip install 'lpu-nn[plot] @ git+https://github.com/akivajp/lpu-nn.git'

Usage

The trainer takes a working directory and a training corpus. The corpus is either a TSV file (source and target in two columns) or one file per column.

$ lpu-nn-train-seq2seq workdir train.tsv --dev-files dev.tsv --test-files test.tsv --gpu 0

It trains a SentencePiece tokenizer, builds the dataset, and writes a checkpoint directory for every metric it improves on (record.best_dev_loss, record.best_dev_bleu, ...), each holding the model, the optimizer state, the configuration and the scores.

Decoding reads from the standard input and writes to the standard output:

$ lpu-nn-run-seq2seq workdir/record.best_dev_loss --gpu 0 < test.txt > hyp.txt

Sequence matching and ranking

The match ranker scores a pair of sequences. Its corpus is a TSV file of three columns: the two sequences and the target score.

$ lpu-nn-train-match-ranker workdir match-train.tsv --dev-files match-dev.tsv --gpu 0

--match-pooler-type selects the architecture: re2 (Yang et al., 2019) or compare-aggregate (Wang and Jiang, 2017). --loss-method selects how the target is used: point for regression on the score, pair for a pairwise ranking loss, classify for a label distribution. The checkpoints are written per ranking metric (record.best_dev_mrr, record.best_dev_map, ...).

Scoring reads pairs from the standard input, one per line:

$ lpu-nn-run-match-ranker workdir/record.best_dev_mrr --gpu 0 < pairs.tsv

--evaluate reports MRR, MAP and recall at k on a labelled corpus instead, and --replies ranks a whole candidate file against each query.

BERT

Pre-training takes a TSV file of sentence pairs and learns a masked language model together with next-sentence prediction.

$ lpu-nn-train-bert workdir train.tsv --dev-files dev.tsv --gpu 0

--universal uses a Universal Transformer (with an adaptive number of steps and a ponder cost) instead of a fixed stack, and --num-token-types 2 adds the segment embedding that distinguishes the two sides of a pair.

Fine-tuning starts from a pre-trained checkpoint. The classifier takes a TSV file of a sentence and its label; the ranker takes a TSV file of pairs.

$ lpu-nn-train-bert-classifier workdir class-train.tsv --dev-files class-dev.tsv \
    --pre-trained-model bert-workdir/record.best_dev_loss \
    --sentencepiece bert-workdir/sp.model --gpu 0
$ lpu-nn-run-bert-classifier workdir/record.best_dev_acc < sentences.txt

--sentencepiece is required alongside --pre-trained-model: each work directory trains its own tokenizer, and fine-tuning reuses the pre-trained embedding, so the two vocabularies have to be the same one. The command refuses to start when they differ rather than writing a checkpoint that cannot be loaded back.

The scorer writes one predicted label per line. --ranking reads sentence<TAB>label instead and reports MRR and precision at k over the known labels. The pair ranker's scorer, lpu-nn-run-bert-ranker, reads sentence1|||sentence2 and writes one score per line, or ranks a candidate file against each query with --replies.

Sequence tagging

The tagger takes a TSV file of a sentence and one tag per token, in the BIO scheme (O, B-LABEL, I-LABEL).

$ lpu-nn-train-tagger workdir tag-train.tsv --dev-files tag-dev.tsv --gpu 0

--encoder-type selects lstm (bidirectional by default), transformer or bert, and --decoder-type selects linear or crf. With bert, pass --pre-trained-model and --sentencepiece as for the other fine-tuning commands. Each evaluation writes the tagged development set to record.latest/pred_dev.txt and reports entity precision, recall and F1, both with and without matching the labels.

Resuming a run

--resume latest picks the training up from the checkpoint in the work directory. The model is rebuilt from the configuration it was saved with and the weights are loaded into it, so the parameters that decide its structure (--embed-size, --hidden-size, --num-layers, ...) keep the values the checkpoint carries; passing a different one reports what it ignored rather than failing to load the weights.

Everything else follows the command line, which is what continued training needs: the corpus, --num-epochs, --batch-size, --optimizer, --learning-rate, --dropout-ratio and the rest of the training settings can all be replaced on a resume.

$ lpu-nn-train-seq2seq workdir more-data.tsv --resume latest \
    --num-epochs 20 --batch-size 64 --optimizer adam

--override-model-params lifts the restriction for the cases where it is safe, such as raising --max-length. A change that alters the shape of a weight still cannot load, and the command says so.

Note that resuming without raising --num-epochs past the epoch already reached does nothing at all: there is no epoch left to run, so no checkpoint is written.

Language modeling and the tokenizer

The language model trains on plain text, one sentence per line, and learns to predict the next token in both directions.

$ lpu-nn-train-embedding workdir corpus.txt --dev-files dev.txt --gpu 0

lpu-nn-run-tokenizer applies a SentencePiece model that any of these commands trained, reading from the standard input:

$ lpu-nn-run-tokenizer workdir/sp.model < text.txt
$ lpu-nn-run-tokenizer workdir/sp.model --format id < text.txt

Serving a sequence-to-sequence model

$ pip install 'lpu-nn[serve]'
$ lpu-nn-serve-seq2seq ja-en=workdir/record.best_dev_bleu --port 8000

It answers GET / with a page for trying the model out, /api/models with the names it was given, and /api/decode with the decoded output as JSON. Several name=path pairs can be served at once.

The server listens on 127.0.0.1 unless --host says otherwise, and it does not run in bottle's debug mode, which would return tracebacks to whoever called it.

Run any command with --help for the full list of options.

Layout

The configuration, logging, progress display and file utilities come from lpu, so they are not duplicated here.

License

MIT, except for the bundled third-party optimizers; see LICENSE and licenses/NOTICE.md.