unitoken 0.1.5

Fast BPE tokenizer/trainer with a Rust core and Python bindings
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

use std::collections::BTreeMap;

use crate::{MyResult, spec::Spec};

use super::{Merge, Word};

/// An immutable BPE model produced by validating a trainer snapshot.
#[derive(Debug)]
pub struct BpeModel<C, I> {
  special_tokens: Vec<String>,
  vocab: BTreeMap<I, Word<C>>,
  merges: Vec<Merge<C, I>>,
}

impl<C, I> BpeModel<C, I> {
  pub(crate) fn new(
    special_tokens: Vec<String>,
    vocab: BTreeMap<I, Word<C>>,
    merges: Vec<Merge<C, I>>,
  ) -> Self {
    Self {
      special_tokens,
      vocab,
      merges,
    }
  }

  /// Reserved special tokens in vocabulary order.
  pub fn special_tokens(&self) -> &[String] {
    &self.special_tokens
  }

  /// Validated token-id vocabulary.
  pub fn vocab(&self) -> &BTreeMap<I, Word<C>> {
    &self.vocab
  }

  /// Validated merge rules in rank order.
  pub fn merges(&self) -> &[Merge<C, I>] {
    &self.merges
  }

  /// Serialize the vocabulary to JSON using `spec`.
  pub fn save_vocab_json<W: std::io::Write>(&self, spec: &dyn Spec<C, I>, mut writer: W) -> MyResult<()> {
    spec.encode_vocab(&mut writer, &self.vocab)
  }

  /// Serialize the merge list to text using `spec`.
  pub fn save_merges_txt<W: std::io::Write>(&self, spec: &dyn Spec<C, I>, mut writer: W) -> MyResult<()> {
    spec.encode_merges(&mut writer, &self.merges)
  }
}