use std::path::Path;
use std::str::FromStr;
use serde_json::Value;
use lindera::dictionary::{Dictionary, UserDictionary};
use lindera::mode::Mode;
use lindera::segmenter::Segmenter;
use lindera::tokenizer::{Tokenizer, TokenizerBuilder};
use crate::error::CoreResult;
use crate::token::TokenView;
pub struct CoreTokenizerBuilder {
inner: TokenizerBuilder,
}
impl CoreTokenizerBuilder {
pub fn new() -> CoreResult<Self> {
Ok(Self {
inner: TokenizerBuilder::new()?,
})
}
pub fn from_file(file_path: &Path) -> CoreResult<Self> {
Ok(Self {
inner: TokenizerBuilder::from_file(file_path)?,
})
}
pub fn set_mode(&mut self, mode: &str) -> CoreResult<&mut Self> {
let mode = Mode::from_str(mode)?;
self.inner.set_segmenter_mode(&mode);
Ok(self)
}
pub fn set_dictionary(&mut self, uri: &str) -> &mut Self {
self.inner.set_segmenter_dictionary(uri);
self
}
pub fn set_user_dictionary(&mut self, uri: &str) -> &mut Self {
self.inner.set_segmenter_user_dictionary(uri);
self
}
pub fn set_keep_whitespace(&mut self, keep_whitespace: bool) -> &mut Self {
self.inner.set_segmenter_keep_whitespace(keep_whitespace);
self
}
pub fn append_character_filter(&mut self, kind: &str, args: &Value) -> &mut Self {
self.inner.append_character_filter(kind, args);
self
}
pub fn append_token_filter(&mut self, kind: &str, args: &Value) -> &mut Self {
self.inner.append_token_filter(kind, args);
self
}
pub fn build(&self) -> CoreResult<CoreTokenizer> {
Ok(CoreTokenizer {
inner: self.inner.build()?,
})
}
}
pub struct CoreTokenizer {
inner: Tokenizer,
}
impl CoreTokenizer {
pub fn from_segmenter(
mode: &str,
dictionary: Dictionary,
user_dictionary: Option<UserDictionary>,
) -> CoreResult<Self> {
let mode = Mode::from_str(mode)?;
let segmenter = Segmenter::new(mode, dictionary, user_dictionary);
Ok(Self {
inner: Tokenizer::new(segmenter),
})
}
pub fn from_tokenizer(tokenizer: Tokenizer) -> Self {
Self { inner: tokenizer }
}
pub fn tokenize(&self, text: &str) -> CoreResult<Vec<TokenView>> {
let tokens = self.inner.tokenize(text)?;
Ok(tokens.into_iter().map(TokenView::from_token).collect())
}
pub fn tokenize_nbest(
&self,
text: &str,
n: usize,
unique: bool,
cost_threshold: Option<i64>,
) -> CoreResult<Vec<(Vec<TokenView>, i64)>> {
let results = self.inner.tokenize_nbest(text, n, unique, cost_threshold)?;
Ok(results
.into_iter()
.map(|(tokens, cost)| {
(
tokens.into_iter().map(TokenView::from_token).collect(),
cost,
)
})
.collect())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn builder_new_succeeds() {
assert!(CoreTokenizerBuilder::new().is_ok());
}
#[test]
fn set_mode_accepts_known_modes() {
let mut builder = CoreTokenizerBuilder::new().expect("builder");
assert!(builder.set_mode("normal").is_ok());
assert!(builder.set_mode("decompose").is_ok());
}
#[test]
fn set_mode_rejects_unknown_mode() {
let mut builder = CoreTokenizerBuilder::new().expect("builder");
assert!(builder.set_mode("definitely-not-a-mode").is_err());
}
#[test]
fn infallible_setters_chain() {
let mut builder = CoreTokenizerBuilder::new().expect("builder");
builder
.set_dictionary("embedded://ipadic")
.set_keep_whitespace(true)
.append_token_filter("japanese_compound_word", &Value::Object(Default::default()));
}
}