Skip to main content

jpreprocess_dictionary/tokenizer/
jpreprocess.rs

1use std::borrow::Cow;
2
3use jpreprocess_core::{
4    error::DictionaryError,
5    token::{Token, Tokenizer},
6    word_entry::WordEntry,
7    JPreprocessResult,
8};
9
10use crate::{
11    dictionary::word_encoding::JPreprocessDictionaryWordEncoding, word_data::get_word_data,
12};
13
14pub struct JPreprocessTokenizer {
15    tokenizer: lindera::tokenizer::Tokenizer,
16}
17
18impl JPreprocessTokenizer {
19    pub fn new(tokenizer: lindera::tokenizer::Tokenizer) -> Self {
20        Self { tokenizer }
21    }
22
23    fn get_word(
24        &self,
25        word_id: lindera_dictionary::viterbi::WordId,
26    ) -> Result<WordEntry, DictionaryError> {
27        if word_id.is_unknown() {
28            Ok(WordEntry::default())
29        } else if word_id.is_system() {
30            Self::get_word_from_prefixdict(
31                &self.tokenizer.segmenter.dictionary.prefix_dictionary,
32                word_id,
33            )
34        } else {
35            let user = &self.tokenizer.segmenter.user_dictionary;
36            user.as_ref()
37                .map_or(Err(DictionaryError::UserDictionaryNotProvided), |user| {
38                    Self::get_word_from_prefixdict(&user.dict, word_id)
39                })
40        }
41    }
42
43    /// PANIC: It must be ensured that the prefix_dict is the correct dictionary for the word_id.
44    pub(super) fn get_word_from_prefixdict(
45        prefix_dict: &lindera_dictionary::dictionary::prefix_dictionary::PrefixDictionary,
46        word_id: lindera_dictionary::viterbi::WordId,
47    ) -> Result<WordEntry, DictionaryError> {
48        if word_id.is_unknown() {
49            Ok(WordEntry::default())
50        } else {
51            let data = get_word_data(
52                &prefix_dict.words_idx_data,
53                &prefix_dict.words_data,
54                Some(word_id.id as usize),
55            )
56            .ok_or(DictionaryError::IdNotFound(word_id.id))?;
57            Ok(JPreprocessDictionaryWordEncoding::deserialize(data)?)
58        }
59    }
60}
61
62impl Tokenizer for JPreprocessTokenizer {
63    fn tokenize<'a>(&'a self, text: &'a str) -> JPreprocessResult<Vec<impl 'a + Token>> {
64        let words = self.tokenizer.tokenize(text).unwrap();
65        words
66            .into_iter()
67            .map(|token| {
68                Ok(JPreprocessToken::new(
69                    token.surface,
70                    self.get_word(token.word_id)?,
71                ))
72            })
73            .collect::<Result<_, _>>()
74    }
75}
76
77pub struct JPreprocessToken<'a> {
78    text: Cow<'a, str>,
79    entry: WordEntry,
80}
81
82impl<'a> JPreprocessToken<'a> {
83    pub(crate) fn new(text: Cow<'a, str>, entry: WordEntry) -> Self {
84        Self { text, entry }
85    }
86}
87
88impl Token for JPreprocessToken<'_> {
89    fn fetch(&mut self) -> Result<(&str, WordEntry), jpreprocess_core::JPreprocessError> {
90        Ok((&self.text, self.entry.clone()))
91    }
92}