use std::path::Path;
use std::str::FromStr;
use std::sync::{Mutex, MutexGuard};
use serde_json::Value;
use lindera::dictionary::{Dictionary, UserDictionary};
use lindera::mode::Mode;
use lindera::segmenter::Segmenter;
use lindera_analysis::tokenizer::{Tokenizer, TokenizerBuilder};
use lindera_analysis::worker::AnalysisWorker;
use crate::error::CoreResult;
use crate::token::{SurfaceView, 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::from_tokenizer(self.inner.build()?))
}
}
pub struct CoreTokenizer {
worker: Mutex<AnalysisWorker>,
}
fn lock_worker(mutex: &Mutex<AnalysisWorker>) -> MutexGuard<'_, AnalysisWorker> {
match mutex.lock() {
Ok(guard) => guard,
Err(poisoned) => {
let mut guard = poisoned.into_inner();
guard.reset();
guard
}
}
}
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::from_tokenizer(Tokenizer::new(segmenter)))
}
pub fn from_tokenizer(tokenizer: Tokenizer) -> Self {
Self {
worker: Mutex::new(tokenizer.into_worker()),
}
}
pub fn tokenize(&self, text: &str) -> CoreResult<Vec<TokenView>> {
let mut worker = lock_worker(&self.worker);
let tokens = worker.tokenize(text)?;
Ok(tokens.into_iter().map(TokenView::from_token).collect())
}
pub fn tokenize_surfaces(&self, text: &str) -> CoreResult<Vec<SurfaceView>> {
let mut worker = lock_worker(&self.worker);
let tokens = worker.tokenize(text)?;
Ok(tokens.into_iter().map(SurfaceView::from_token).collect())
}
pub fn tokenize_nbest(
&self,
text: &str,
n: usize,
unique: bool,
cost_threshold: Option<i64>,
) -> CoreResult<Vec<(Vec<TokenView>, i64)>> {
let mut worker = lock_worker(&self.worker);
let results = worker.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()));
}
#[test]
fn core_tokenizer_is_send_and_sync() {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<CoreTokenizer>();
assert_send_sync::<CoreTokenizerBuilder>();
}
#[cfg(feature = "embed-ipadic")]
mod with_ipadic {
use super::super::*;
fn ipadic_tokenizer() -> CoreTokenizer {
let dictionary = match lindera::dictionary::load_dictionary("embedded://ipadic") {
Ok(dictionary) => dictionary,
Err(err) => panic!("failed to load embedded IPADIC: {err}"),
};
match CoreTokenizer::from_segmenter("normal", dictionary, None) {
Ok(tokenizer) => tokenizer,
Err(err) => panic!("failed to build CoreTokenizer: {err}"),
}
}
#[test]
fn tokenize_repeated_calls_are_stable() {
let tokenizer = ipadic_tokenizer();
let first = match tokenizer.tokenize("すもももももももものうち") {
Ok(tokens) => tokens,
Err(err) => panic!("tokenize failed: {err}"),
};
assert!(!first.is_empty());
for _ in 0..100 {
let again = match tokenizer.tokenize("すもももももももものうち") {
Ok(tokens) => tokens,
Err(err) => panic!("tokenize failed: {err}"),
};
assert_eq!(first.len(), again.len());
for (a, b) in first.iter().zip(again.iter()) {
assert_eq!(a.surface, b.surface);
assert_eq!(a.byte_start, b.byte_start);
assert_eq!(a.byte_end, b.byte_end);
assert_eq!(a.details, b.details);
}
}
}
#[test]
fn tokenize_nbest_repeated_calls_are_stable() {
let tokenizer = ipadic_tokenizer();
let first = match tokenizer.tokenize_nbest("すもももももももものうち", 3, false, None)
{
Ok(results) => results,
Err(err) => panic!("tokenize_nbest failed: {err}"),
};
assert!(!first.is_empty());
for _ in 0..10 {
let again =
match tokenizer.tokenize_nbest("すもももももももものうち", 3, false, None)
{
Ok(results) => results,
Err(err) => panic!("tokenize_nbest failed: {err}"),
};
assert_eq!(first.len(), again.len());
for ((a_tokens, a_cost), (b_tokens, b_cost)) in first.iter().zip(again.iter()) {
assert_eq!(a_cost, b_cost);
assert_eq!(a_tokens.len(), b_tokens.len());
}
}
}
#[test]
fn tokenize_surfaces_matches_tokenize() {
let tokenizer = ipadic_tokenizer();
let texts = [
"すもももももももものうち",
"関西国際空港限定トートバッグ",
"",
];
for _ in 0..3 {
for text in texts {
let full = match tokenizer.tokenize(text) {
Ok(tokens) => tokens,
Err(err) => panic!("tokenize failed: {err}"),
};
let surfaces = match tokenizer.tokenize_surfaces(text) {
Ok(views) => views,
Err(err) => panic!("tokenize_surfaces failed: {err}"),
};
assert_eq!(full.len(), surfaces.len(), "count for {text:?}");
for (a, b) in full.iter().zip(surfaces.iter()) {
assert_eq!(a.surface, b.surface);
assert_eq!(a.byte_start, b.byte_start);
assert_eq!(a.byte_end, b.byte_end);
}
}
}
}
#[test]
fn tokenize_recovers_from_poisoned_lock() {
use std::sync::Arc;
let tokenizer = Arc::new(ipadic_tokenizer());
let expected = match tokenizer.tokenize("すもももももももものうち") {
Ok(tokens) => tokens,
Err(err) => panic!("tokenize failed: {err}"),
};
let poisoner = Arc::clone(&tokenizer);
let result = std::thread::spawn(move || {
let _guard = lock_worker(&poisoner.worker);
panic!("intentional panic to poison the worker lock");
})
.join();
assert!(result.is_err(), "poisoning thread must have panicked");
assert!(tokenizer.worker.is_poisoned(), "lock must be poisoned");
let after = match tokenizer.tokenize("すもももももももものうち") {
Ok(tokens) => tokens,
Err(err) => panic!("tokenize after poison failed: {err}"),
};
assert_eq!(expected.len(), after.len());
for (a, b) in expected.iter().zip(after.iter()) {
assert_eq!(a.surface, b.surface);
}
}
}
}