ct2rs 0.9.19

Rust bindings for OpenNMT/CTranslate2
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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

// hf.rs
//
// Copyright (c) 2023-2024 Junpei Kawamoto
//
// This software is released under the MIT License.
//
// http://opensource.org/licenses/mit-license.php

//! A module for utilizing the tokenizer provided by the Hugging Face's
//! [`tokenizers` crate](https://docs.rs/tokenizers/).
//!
//! This module allows the creation of a [`Tokenizer`] structure instance by specifying
//! the path to a directory containing `tokenizer.json`. The created instance can then be
//! used in [`Translator`](crate::Translator) and [`Generator`](crate::Generator).
//!
//! # Example
//! Here is an example of how to create an instance of the [`Tokenizer`] structure
//! and then use it to create an instance of the [`Translator`](crate::Translator):
//!
//! ```no_run
//! # use anyhow::Result;
//! #
//! use ct2rs::{Config, Translator};
//! use ct2rs::tokenizers::hf::Tokenizer;
//!
//! # fn main() -> Result<()> {
//! let path = "/path/to/model";
//! let t = Translator::with_tokenizer(&path, Tokenizer::new(&path)?, &Config::default())?;
//! # Ok(())
//! # }
//! ```

use std::ops::{Deref, DerefMut};
use std::path::Path;

use anyhow::{anyhow, Result};
use tokenizers::Decoder;

const TOKENIZER_FILENAME: &str = "tokenizer.json";

/// Tokenizer that utilize the tokenizer provided by the Hugging Face's `tokenizers` crate.
pub struct Tokenizer {
    tokenizer: tokenizers::Tokenizer,
    special_token: bool,
}

impl Tokenizer {
    /// Create a tokenizer instance by specifying the path to a directory containing
    /// `tokenizer.json`.
    pub fn new<T: AsRef<Path>>(path: T) -> Result<Self> {
        Tokenizer::from_file(path.as_ref().join(TOKENIZER_FILENAME))
    }

    /// Create a tokenizer instance by specifying the path to `tokenizer.json`.
    pub fn from_file<T: AsRef<Path>>(path: T) -> Result<Self> {
        Ok(Self {
            tokenizer: tokenizers::tokenizer::Tokenizer::from_file(path)
                .map_err(|err| anyhow!("failed to load a tokenizer: {err}"))?,
            special_token: true,
        })
    }

    /// Disable adding special tokens.
    ///
    /// See also [Special tokens in translation](https://opennmt.net/CTranslate2/guides/transformers.html#special-tokens-in-translation).
    pub fn disable_spacial_token(&mut self) -> &mut Self {
        self.special_token = false;
        self
    }

    /// Returns a mutable reference of the `tokenizers::Tokenizer`.
    #[inline]
    pub fn inner(&mut self) -> &mut tokenizers::Tokenizer {
        self.deref_mut()
    }
}

impl crate::Tokenizer for Tokenizer {
    /// Encodes a given string into a sequence of tokens.
    ///
    /// This function takes a reference to a string and returns a vector of token strings
    /// resulting from the tokenization process.
    ///
    /// # Arguments
    /// * `input` - A reference to the string to be tokenized.
    ///
    /// # Returns
    /// A `Result` containing either the vector of tokens if successful or an error if the
    /// tokenization fails.
    fn encode(&self, input: &str) -> Result<Vec<String>> {
        self.tokenizer
            .encode(input, self.special_token)
            .map(|r| r.get_tokens().to_vec())
            .map_err(|err| anyhow!("failed to encode the given input: {err}"))
    }

    /// Decodes a given sequence of tokens back into a single string.
    ///
    /// This function takes a vector of token strings and reconstructs the original string.
    ///
    /// # Arguments
    /// * `tokens` - A vector of strings representing the tokens to be decoded.
    ///
    /// # Returns
    /// A `Result` containing either the reconstructed string if successful or an error if the
    /// decoding fails.
    fn decode(&self, tokens: Vec<String>) -> Result<String> {
        let decoder = self
            .tokenizer
            .get_decoder()
            .ok_or_else(|| anyhow!("no decoder is provided"))?;

        decoder
            .decode(tokens)
            .map_err(|err| anyhow!("failed to decode: {err}"))
    }
}

impl From<tokenizers::Tokenizer> for Tokenizer {
    fn from(tokenizer: tokenizers::Tokenizer) -> Self {
        Self {
            tokenizer,
            special_token: true,
        }
    }
}

impl Deref for Tokenizer {
    type Target = tokenizers::Tokenizer;

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.tokenizer
    }
}

impl DerefMut for Tokenizer {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.tokenizer
    }
}