BpeTokenizer.train constructor

BpeTokenizer.train(String corpus, { required int targetVocabSize, int minCount = 2, })

Train a fresh tokenizer from a text corpus.

targetVocabSize must be >= 256; training stops when the vocab reaches that size or when no more valid pairs exist. minCount filters out pair candidates that occur fewer than this many times — a small speedup / stability guard for tiny corpora.

Implementation

factory BpeTokenizer.train(
  String corpus, {
  required int targetVocabSize,
  int minCount = 2,
}) {
  if (targetVocabSize < 256) {
    throw ArgumentError(
      'BpeTokenizer.train: targetVocabSize must be >= 256; got $targetVocabSize',
    );
  }
  final bytes = utf8.encode(corpus);
  if (bytes.isEmpty) {
    throw ArgumentError('BpeTokenizer.train: corpus is empty');
  }
  // Working sequence — each entry is a current token id.
  var seq = List<int>.of(bytes);
  final merges = <List<int>>[];
  var nextId = 256;

  while (nextId < targetVocabSize) {
    // Count adjacent pairs.
    final counts = <int, int>{};
    for (int i = 0; i + 1 < seq.length; i++) {
      final k = _pairKey(seq[i], seq[i + 1]);
      counts[k] = (counts[k] ?? 0) + 1;
    }
    if (counts.isEmpty) break;

    // Pick the most frequent pair (ties broken by lower key for
    // determinism across runs).
    int bestKey = -1;
    int bestCount = 0;
    counts.forEach((k, c) {
      if (c > bestCount ||
          (c == bestCount && (bestKey == -1 || k < bestKey))) {
        bestCount = c;
        bestKey = k;
      }
    });
    if (bestCount < minCount) break;

    final a = bestKey >> 20;
    final b = bestKey & 0xFFFFF;
    merges.add([a, b]);

    // Apply the merge to `seq` in a single pass.
    final merged = <int>[];
    int i = 0;
    while (i < seq.length) {
      if (i + 1 < seq.length && seq[i] == a && seq[i + 1] == b) {
        merged.add(nextId);
        i += 2;
      } else {
        merged.add(seq[i]);
        i += 1;
      }
    }
    seq = merged;
    nextId += 1;
  }

  return BpeTokenizer._(nextId, merges);
}