use crate::error::TokenizerError;
use std::collections::HashMap;
use std::fs::File;
use std::hash::Hash;
use std::io::{BufRead, BufReader};
pub(crate) fn swap_key_values<T: Clone, U: Hash + Eq + Copy>(
input_hashmap: &HashMap<T, U>,
) -> HashMap<U, T> {
input_hashmap
.iter()
.map(|(key, &value)| (value, key.clone()))
.collect()
}
pub trait Vocab {
fn unknown_value() -> &'static str;
fn get_unknown_value(&self) -> &'static str;
fn values(&self) -> &HashMap<String, i64>;
fn indices(&self) -> &HashMap<i64, String>;
fn special_values(&self) -> &HashMap<String, i64>;
fn special_indices(&self) -> &HashMap<i64, String>;
fn from_file(path: &str) -> Result<Self, TokenizerError>
where
Self: std::marker::Sized;
fn read_vocab_file(path: &str) -> Result<HashMap<String, i64>, TokenizerError> {
let f = File::open(path).map_err(|e| {
TokenizerError::FileNotFound(format!("{} vocabulary file not found :{}", path, e))
})?;
let br = BufReader::new(f);
let mut data = HashMap::new();
for (index, line) in br.lines().enumerate() {
let line = match line {
Ok(value) => value,
Err(e) => {
return Err(TokenizerError::VocabularyParsingError(e.to_string()));
}
};
data.insert(line.trim().to_owned(), index as i64);
}
Ok(data)
}
fn _token_to_id(
&self,
token: &str,
values: &HashMap<String, i64>,
special_values: &HashMap<String, i64>,
unknown_value: &str,
) -> i64 {
match special_values.get(token) {
Some(index) => *index,
None => match values.get(token) {
Some(index) => *index,
None => *values.get(unknown_value).unwrap(),
},
}
}
fn _id_to_token(
&self,
id: &i64,
indices: &HashMap<i64, String>,
special_indices: &HashMap<i64, String>,
unknown_value: &str,
) -> String {
match special_indices.get(id) {
Some(token) => token.clone(),
None => match indices.get(id) {
Some(token) => token.clone(),
None => unknown_value.to_owned(),
},
}
}
fn _register_as_special_value(
token: &str,
values: &HashMap<String, i64>,
special_values: &mut HashMap<String, i64>,
) -> Result<(), TokenizerError> {
let token_id = match values.get(token) {
Some(index) => *index,
None => {
return Err(TokenizerError::TokenNotFound(format!(
"The special value {} could not be found in the vocabulary",
token
)));
}
};
special_values.insert(String::from(token), token_id);
Ok(())
}
fn token_to_id(&self, token: &str) -> i64;
fn id_to_token(&self, id: &i64) -> String;
fn convert_tokens_to_ids(&self, tokens: &[&str]) -> Vec<i64> {
tokens.iter().map(|v| self.token_to_id(v)).collect()
}
}
#[derive(Debug, Clone)]
pub struct BaseVocab {
pub values: HashMap<String, i64>,
pub indices: HashMap<i64, String>,
pub unknown_value: &'static str,
pub special_values: HashMap<String, i64>,
pub special_indices: HashMap<i64, String>,
}
impl Vocab for BaseVocab {
fn unknown_value() -> &'static str {
"[UNK]"
}
fn get_unknown_value(&self) -> &'static str {
"[UNK]"
}
fn values(&self) -> &HashMap<String, i64> {
&self.values
}
fn indices(&self) -> &HashMap<i64, String> {
&self.indices
}
fn special_values(&self) -> &HashMap<String, i64> {
&self.special_values
}
fn special_indices(&self) -> &HashMap<i64, String> {
&self.special_indices
}
fn from_file(path: &str) -> Result<BaseVocab, TokenizerError> {
let values = BaseVocab::read_vocab_file(path)?;
let mut special_values = HashMap::new();
let unknown_value = BaseVocab::unknown_value();
BaseVocab::_register_as_special_value(unknown_value, &values, &mut special_values)?;
let indices = swap_key_values(&values);
let special_indices = swap_key_values(&special_values);
Ok(BaseVocab {
values,
indices,
unknown_value,
special_values,
special_indices,
})
}
fn token_to_id(&self, token: &str) -> i64 {
self._token_to_id(
token,
&self.values,
&self.special_values,
&self.unknown_value,
)
}
fn id_to_token(&self, id: &i64) -> String {
self._id_to_token(
&id,
&self.indices,
&self.special_indices,
&self.unknown_value,
)
}
}
#[cfg(test)]
mod tests {
extern crate anyhow;
use super::*;
use std::io::Write;
#[test]
fn test_create_object() {
let values: HashMap<String, i64> = HashMap::new();
let special_values: HashMap<String, i64> = HashMap::new();
let indices: HashMap<i64, String> = HashMap::new();
let special_indices: HashMap<i64, String> = HashMap::new();
let unknown_value = BaseVocab::unknown_value();
let base_vocab = BaseVocab {
values,
indices,
unknown_value,
special_values,
special_indices,
};
assert_eq!(base_vocab.unknown_value, "[UNK]");
assert_eq!(base_vocab.unknown_value, BaseVocab::unknown_value());
assert_eq!(base_vocab.values, *base_vocab.values());
assert_eq!(base_vocab.special_values, *base_vocab.special_values());
}
#[test]
fn test_create_object_from_file() -> anyhow::Result<()> {
let mut vocab_file = tempfile::NamedTempFile::new()?;
write!(vocab_file, "hello \n world \n [UNK] \n !")?;
let path = vocab_file.into_temp_path();
let target_values: HashMap<String, i64> = [
("hello".to_owned(), 0),
("world".to_owned(), 1),
("[UNK]".to_owned(), 2),
("!".to_owned(), 3),
]
.iter()
.cloned()
.collect();
let special_values: HashMap<String, i64> =
[("[UNK]".to_owned(), 2)].iter().cloned().collect();
let base_vocab = BaseVocab::from_file(path.to_path_buf().to_str().unwrap())?;
assert_eq!(base_vocab.unknown_value, "[UNK]");
assert_eq!(base_vocab.values, target_values);
assert_eq!(base_vocab.special_values, special_values);
drop(path);
Ok(())
}
#[test]
#[should_panic]
fn test_create_object_from_file_without_unknown_token() {
let mut vocab_file = tempfile::NamedTempFile::new().unwrap();
write!(vocab_file, "hello \n world \n !").unwrap();
let path = vocab_file.into_temp_path();
let _base_vocab = BaseVocab::from_file(path.to_path_buf().to_str().unwrap()).unwrap();
}
#[test]
fn test_encode_tokens() -> anyhow::Result<()> {
let mut vocab_file = tempfile::NamedTempFile::new()?;
write!(vocab_file, "hello \n world \n [UNK] \n !")?;
let path = vocab_file.into_temp_path();
let base_vocab = BaseVocab::from_file(path.to_path_buf().to_str().unwrap())?;
assert_eq!(base_vocab.token_to_id("hello"), 0);
assert_eq!(base_vocab.token_to_id("world"), 1);
assert_eq!(base_vocab.token_to_id("!"), 3);
assert_eq!(base_vocab.token_to_id("[UNK]"), 2);
assert_eq!(base_vocab.token_to_id("oov_value"), 2);
drop(path);
Ok(())
}
#[test]
fn test_decode_tokens() -> anyhow::Result<()> {
let mut vocab_file = tempfile::NamedTempFile::new()?;
write!(vocab_file, "hello \n world \n [UNK] \n !")?;
let path = vocab_file.into_temp_path();
let base_vocab = BaseVocab::from_file(path.to_path_buf().to_str().unwrap())?;
assert_eq!(base_vocab.id_to_token(&(0_i64)), "hello");
assert_eq!(base_vocab.id_to_token(&(1_i64)), "world");
assert_eq!(base_vocab.id_to_token(&(3_i64)), "!");
assert_eq!(base_vocab.id_to_token(&(2_i64)), "[UNK]");
drop(path);
Ok(())
}
}