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::space_penalty::SpacePenaltyConfig;
use lindera_analysis::tokenizer::{Tokenizer, TokenizerBuilder};
use lindera_analysis::worker::AnalysisWorker;
use crate::error::{CoreError, CoreResult};
use crate::token::{SurfaceView, TokenView};
#[derive(Debug, Clone, PartialEq)]
enum SpacePenaltySetting {
Default,
Off,
FromDictionary,
Rules(SpacePenaltyConfig),
}
impl SpacePenaltySetting {
fn from_value(value: &Value) -> CoreResult<Self> {
match value {
Value::Null => Ok(Self::Default),
Value::Bool(false) => Ok(Self::Off),
Value::Bool(true) => Ok(Self::FromDictionary),
Value::Object(_) => serde_json::from_value(value.clone())
.map(Self::Rules)
.map_err(|err| {
CoreError::invalid_argument(format!("invalid space_penalty rules: {err}"))
}),
other => Err(CoreError::invalid_argument(format!(
"space_penalty must be null, a boolean or an object with \"rules\", got {other}"
))),
}
}
fn apply(&self, segmenter: &mut Segmenter) -> CoreResult<()> {
match self {
Self::Default => {}
Self::Off => segmenter.set_space_penalty(None)?,
Self::FromDictionary => segmenter.set_space_penalty_from_dictionary()?,
Self::Rules(config) => segmenter.set_space_penalty(Some(config.clone()))?,
}
Ok(())
}
}
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 set_space_penalty(&mut self, value: &Value) -> CoreResult<&mut Self> {
match SpacePenaltySetting::from_value(value)? {
SpacePenaltySetting::Default => {
self.inner.reset_segmenter_space_penalty();
}
SpacePenaltySetting::Off => {
self.inner.set_segmenter_space_penalty(None);
}
SpacePenaltySetting::FromDictionary => {
self.inner.set_segmenter_space_penalty_from_dictionary(true);
}
SpacePenaltySetting::Rules(config) => {
self.inner.set_segmenter_space_penalty(Some(&config));
}
}
Ok(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_segmenter_with_space_penalty(
mode: &str,
dictionary: Dictionary,
user_dictionary: Option<UserDictionary>,
space_penalty: &Value,
) -> CoreResult<Self> {
let setting = SpacePenaltySetting::from_value(space_penalty)?;
let mode = Mode::from_str(mode)?;
let mut segmenter = Segmenter::new(mode, dictionary, user_dictionary);
setting.apply(&mut segmenter)?;
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 set_space_penalty_accepts_the_config_forms() {
let mut builder = CoreTokenizerBuilder::new().expect("builder");
for value in [
Value::Null,
Value::Bool(true),
Value::Bool(false),
serde_json::json!({ "rules": [{ "pos": ["JKS"], "cost": 6000 }] }),
serde_json::json!({ "rules": [] }),
] {
assert!(
builder.set_space_penalty(&value).is_ok(),
"rejected {value}"
);
}
}
#[test]
fn set_space_penalty_rejects_other_values() {
let mut builder = CoreTokenizerBuilder::new().expect("builder");
for value in [
serde_json::json!(1),
serde_json::json!("false"),
serde_json::json!(["JKS"]),
serde_json::json!({ "rules": "JKS" }),
serde_json::json!({ "rules": [{ "pos": ["JKS"] }] }),
] {
let err = match builder.set_space_penalty(&value) {
Ok(_) => panic!("accepted {value}"),
Err(err) => err,
};
assert_eq!(err.kind(), crate::error::ErrorKind::InvalidArgument);
}
}
#[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);
}
}
#[test]
fn space_penalty_true_fails_on_a_dictionary_without_rules() {
for (value, should_build) in [
(Value::Null, true),
(Value::Bool(false), true),
(Value::Bool(true), false),
] {
let mut builder = CoreTokenizerBuilder::new().expect("builder");
builder
.set_dictionary("embedded://ipadic")
.set_space_penalty(&value)
.expect("valid value");
assert_eq!(builder.build().is_ok(), should_build, "value {value}");
}
}
#[test]
fn from_segmenter_with_space_penalty_on_ipadic() {
for (value, should_build) in [
(Value::Null, true),
(Value::Bool(false), true),
(Value::Bool(true), false),
(
serde_json::json!({ "rules": [{ "pos": ["助詞"], "cost": 1000 }] }),
true,
),
] {
let dictionary = match lindera::dictionary::load_dictionary("embedded://ipadic") {
Ok(dictionary) => dictionary,
Err(err) => panic!("failed to load embedded IPADIC: {err}"),
};
let result = CoreTokenizer::from_segmenter_with_space_penalty(
"normal", dictionary, None, &value,
);
assert_eq!(result.is_ok(), should_build, "value {value}");
}
}
}
#[cfg(feature = "embed-ko-dic")]
mod with_ko_dic {
use super::super::*;
fn tags(space_penalty: &Value, text: &str) -> Vec<(String, String)> {
let mut builder = CoreTokenizerBuilder::new().expect("builder");
builder
.set_dictionary("embedded://ko-dic")
.set_space_penalty(space_penalty)
.expect("valid value");
let tokenizer = match builder.build() {
Ok(tokenizer) => tokenizer,
Err(err) => panic!("build failed: {err}"),
};
match tokenizer.tokenize(text) {
Ok(tokens) => tokens
.into_iter()
.map(|t| (t.surface, t.details.first().cloned().unwrap_or_default()))
.collect(),
Err(err) => panic!("tokenize failed: {err}"),
}
}
#[test]
fn space_penalty_setting_changes_ko_dic_output() {
let text = "서울 시 에서 출발";
let default = tags(&Value::Null, text);
let off = tags(&Value::Bool(false), text);
let required = tags(&Value::Bool(true), text);
let shipped = serde_json::json!({ "rules": [
{ "pos": ["EC", "EF", "EP", "ETM", "ETN", "VCP", "XSA", "XSN", "XSV"], "cost": 3000 },
{ "pos": ["JC", "JKB", "JKC", "JKG", "JKO", "JKQ", "JKS", "JKV", "JX"], "cost": 6000 }
] });
let explicit = tags(&shipped, text);
assert_ne!(
default, off,
"turning the penalty off must change the output"
);
assert_eq!(default, required);
assert_eq!(default, explicit);
assert!(
off.iter().any(|(s, tag)| s == "시" && tag == "EP"),
"{off:?}"
);
assert!(
!default.iter().any(|(s, tag)| s == "시" && tag == "EP"),
"{default:?}"
);
}
#[test]
fn from_segmenter_with_space_penalty_matches_the_builder() {
let text = "서울 시 에서 출발";
for value in [Value::Null, Value::Bool(false), Value::Bool(true)] {
let dictionary = match lindera::dictionary::load_dictionary("embedded://ko-dic") {
Ok(dictionary) => dictionary,
Err(err) => panic!("failed to load embedded ko-dic: {err}"),
};
let tokenizer = match CoreTokenizer::from_segmenter_with_space_penalty(
"normal", dictionary, None, &value,
) {
Ok(tokenizer) => tokenizer,
Err(err) => panic!("build failed for {value}: {err}"),
};
let from_instance: Vec<(String, String)> = match tokenizer.tokenize(text) {
Ok(tokens) => tokens
.into_iter()
.map(|t| (t.surface, t.details.first().cloned().unwrap_or_default()))
.collect(),
Err(err) => panic!("tokenize failed: {err}"),
};
assert_eq!(from_instance, tags(&value, text), "value {value}");
}
}
}
}