| 18 | } |
| 19 | |
| 20 | pub fn encode(&self, text: &str, include_bos: bool, include_eos: bool) -> Vec<i64> { |
| 21 | let pre = if include_bos { |
| 22 | vec![BOS_TOKEN_ID] |
| 23 | } else { |
| 24 | vec![] |
| 25 | }; |
| 26 | |
| 27 | let post = if include_eos { |
| 28 | vec![EOS_TOKEN_ID] |
| 29 | } else { |
| 30 | vec![] |
| 31 | }; |
| 32 | |
| 33 | let token_ids = self.spm.encode(text, None, std::usize::MAX, &TruncationStrategy::LongestFirst, 0).token_ids; |
| 34 | |
| 35 | [pre, token_ids, post] |
| 36 | .into_iter() |
| 37 | .flat_map(|v| v.into_iter()) |
| 38 | .collect() |
| 39 | } |
| 40 | |
| 41 | pub fn decode(&self, tokens: &[i64], skip_special_tokens: bool) -> String { |
| 42 | let clean_spaces = false; |