mojolearn.lm_corpus#

Tokenize a training corpus once and train a language model on the ids.

The tokenizer is part of the model: it is chosen before the first training step, and the model’s ids mean something only through it. This module prepares a corpus with a BPE vocabulary, caches the token ids, serves training batches, and checks that a tokenizer is the one a model was trained on.

  • With no vocab, prepare() trains a byte-level BPE vocabulary on the corpus’s training range (never validation or test data), tokenizes the whole corpus once and caches the ids.

  • With vocab=, your own token map is used and nothing is trained: a rank file, an (encoder.json, vocab.bpe) pair, a BpeTokenizer or a TrainedBpeVocabulary.

Cache#

Files are written under cache_dir (default $MOJOLEARN_LM_CACHE, else ~/.cache/mojolearn/lm):

vocab/<corpus sha16>-v<vocab_size>-f<min_frequency>-s<sample bytes>/
    ranks.tsv, vocabulary.json            (a trained vocabulary)
vocab/user-<vocabulary sha16>/ranks.tsv, vocabulary.json   (yours)
tokens/<corpus sha16>-<vocabulary sha16>-d<document bytes>/
    tokens.i32, manifest.json

A rerun with the same corpus and vocabulary reuses the cache after re-hashing it against its manifest; a mismatch raises ValueError rather than silently rebuilding. Nothing is written into the package.

Vocabulary identity#

TokenBatches.data_schedule() records the vocabulary’s identity (the SHA-256 of its canonical rank file and its n_vocab). The trainer refuses a schedule whose n_vocab differs from its vocab_size and keeps the schedule in every checkpoint, and tokenizer_for() only returns a tokenizer that matches it, so ids are never decoded with the wrong table.

Documents#

BPE never merges across a pre-token boundary, and where a long file is cut changes the ids at every cut, so the cut is fixed and recorded: a document never spans the corpus manifest’s train / validation / test boundaries, is at most document_bytes long, and ends at the last newline at or before that limit (or exactly at the limit when there is none). Each document is encoded alone and the ids are concatenated. <|endoftext|> is neither inserted nor recognized in the text.

Examples

>>> from mojolearn import lm_corpus, LanguageModelConfig, LanguageModelTrainer
>>> corpus = lm_corpus.prepare("corpus/input.txt")               # train a vocabulary
>>> corpus = lm_corpus.prepare("corpus/input.txt", vocab="ranks.tsv")  # use yours
>>> batches = corpus.batches(batch=1, length=2048)
>>> schedule = batches.data_schedule()
>>> ids = batches.ids(0)
>>> tok = lm_corpus.tokenizer_for(schedule, corpus.vocabulary_path)

Functions

prepare(corpus, *[, vocab, cache_dir, ...])

Tokenize a corpus file once and return a TokenizedCorpus.

require_vocabulary(model, tokenizer)

Check that tokenizer is the vocabulary model was trained on.

tokenizer_for(model, vocabulary)

A tokenizer for vocabulary, checked against the vocabulary model was trained on.

Classes

TokenizedCorpus(tokens_dir, vocabulary_path)

A prepared corpus: cached token ids and the vocabulary they use.

TokenBatches(tokens_dir, batch, length)

Training batches from a tokenized corpus.
mojolearn.lm_corpus.prepare(corpus, *, vocab=None, cache_dir=None, vocab_size=50256, min_frequency=2, vocab_sample_bytes=10000000, document_bytes=1048576, batch_documents=16, progress=None)[source]#

Tokenize a corpus file once and return a TokenizedCorpus.

If a manifest.json with schema mojolearn.byte-lm.corpus.v1 sits beside the corpus, the file is checked against its SHA-256 and length and its train / validation / test ranges are used; otherwise the whole file is the training range. Results are cached (see the module documentation).

Parameters:

  • corpus (str or path-like) – Path to the corpus text file.

  • vocab (str, path-like, tuple, BpeTokenizer, TrainedBpeVocabulary or None, default None) – None trains a vocabulary of vocab_size ranks (n_vocab = vocab_size + 1 with <|endoftext|>) on the first vocab_sample_bytes of the training range. Otherwise your token map: a rank file, an (encoder.json, vocab.bpe) pair, a BpeTokenizer or a TrainedBpeVocabulary; nothing is trained. "bytes" raises ValueError (byte-level training needs no vocabulary).

  • cache_dir (str, path-like or None, default None) – Cache location; None uses default_cache_dir().

  • vocab_size (int, default 50256) – Ranks to train (GPT-2’s vocabulary size, without <|endoftext|>).

  • min_frequency (int, default 2) – Minimum pair frequency for a merge.

  • vocab_sample_bytes (int, default 10_000_000) – Bytes of the training range used to train the vocabulary.

  • document_bytes (int, default 1048576) – Maximum document length; see the module documentation.

  • batch_documents (int, default 16) – Documents tokenized per batch.

  • progress (callable or None, default None) – Optional progress callback.

Returns:

TokenizedCorpus

Raises:

ValueError – For a non-positive size argument, a corpus that does not match its manifest, or a cache that does not match its manifest.

class mojolearn.lm_corpus.TokenizedCorpus(tokens_dir, vocabulary_path)[source]#

Bases: object

A prepared corpus: cached token ids and the vocabulary they use.

Returned by prepare().

Variables:

  • tokens_dir (pathlib.Path) – Directory holding tokens.i32 (the int32 ids) and manifest.json.

  • vocabulary_path (pathlib.Path) – The vocabulary’s canonical rank file.

  • manifest (dict) – The parsed token manifest.

property vocabulary#

The vocabulary identity record (schema, sha256, n_vocab).

property n_vocab#

Number of token ids in the vocabulary, <|endoftext|> included.

tokenizer()[source]#

A BpeTokenizer for this corpus’s vocabulary.

batches(batch, length)[source]#

Training batches over this corpus; see TokenBatches.

class mojolearn.lm_corpus.TokenBatches(tokens_dir, batch, length)[source]#

Bases: object

Training batches from a tokenized corpus.

Step k (zero-based), row b reads the length + 1 ids starting at lo + (k * batch * length + b * length) % (hi - lo - length - 1), where [lo, hi) is the training range in tokens, so a run of any length reads no validation or test id. Targets are the inputs shifted by one id. The token file is verified against its manifest when opened.

Parameters:

  • tokens_dir (str or path-like) – A TokenizedCorpus tokens_dir.

  • batch (int) – Rows per batch.

  • length (int) – Tokens per row (each row holds length + 1 ids).

Raises:

ValueError – If the token file does not match its manifest or the training range is shorter than one row.

property vocabulary#

The vocabulary identity record (schema, sha256, n_vocab).

ids(step_index)[source]#

The batch for training step step_index.

Returns:

Array of shape (batch, length + 1), int32

prefetch(start_step, steps, *, depth=2)[source]#

Iterate over (step, ids) for steps steps from start_step.

Each batch is made when the consumer asks for it, and an error is raised at its own step, before any later batch is made.

Parameters:

  • start_step (int) – First step, at least 0.

  • steps (int) – Number of steps, at least 0.

  • depth (int, default 2) – Maximum number of batches alive at once; at least 1.

Returns:

iterator of (int, Array)

describe()[source]#

A dict describing the token file, its hashes, the training range and the schedule.

data_schedule(**extra)[source]#

The trainer’s data_schedule: which ids, in which order, from which vocabulary.

Pass it to the language model trainer, which keeps it in every checkpoint. Keyword arguments are added to the record.

Returns:

dict

mojolearn.lm_corpus.require_vocabulary(model, tokenizer)[source]#

Check that tokenizer is the vocabulary model was trained on.

Parameters:

  • model (trainer, dict) – A trainer, a checkpoint state, a data schedule or a tokens manifest.

  • tokenizer (BpeTokenizer) – The tokenizer to check.

Returns:

BpeTokenizer – tokenizer, unchanged.

Raises:

ValueError – If model carries no vocabulary identity, or the tokenizer’s SHA-256 or n_vocab differs from it.

mojolearn.lm_corpus.tokenizer_for(model, vocabulary)[source]#

A tokenizer for vocabulary, checked against the vocabulary model was trained on.

Parameters:

  • model (trainer, dict) – A trainer, a checkpoint state, a data schedule or a tokens manifest.

  • vocabulary (str, path-like, tuple, BpeTokenizer or TrainedBpeVocabulary) – Any form prepare() accepts as vocab.

Returns:

BpeTokenizer

Raises:

ValueError – If the vocabulary does not match the model’s (see require_vocabulary()).