Skip to main content

sudachi/
text_normalizer.rs

1/*
2 * Copyright (c) 2026 Works Applications Co., Ltd.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use std::sync::Arc;
18
19use crate::config::ConfigBuilder;
20use crate::dic::grammar::Grammar;
21use crate::dic::DictionaryAccess;
22use crate::input_text::InputBuffer;
23use crate::plugin::input_text::default_input_text::DefaultInputTextPlugin;
24use crate::plugin::input_text::InputTextPlugin;
25use crate::prelude::*;
26
27/// Applies input-text normalization used by tokenizer input processing.
28///
29/// By default, this uses `DefaultInputTextPlugin`. When built from a dictionary,
30/// it applies that dictionary's configured input-text plugins.
31pub struct TextNormalizer<D = DefaultInputTextPlugin> {
32    source: D,
33    input: InputBuffer,
34}
35
36impl TextNormalizer<DefaultInputTextPlugin> {
37    /// Create a text normalizer using the default input-text plugin.
38    pub fn new(grammar: &Grammar) -> SudachiResult<Self> {
39        Ok(Self {
40            source: set_up_default_plugin(grammar)?,
41            input: InputBuffer::new(),
42        })
43    }
44
45    /// Create a text normalizer using the default input-text plugin and an empty grammar.
46    pub fn try_default() -> SudachiResult<Self> {
47        let grammar = Grammar::empty();
48        Self::new(&grammar)
49    }
50
51    pub fn normalize(&mut self, text: &str) -> SudachiResult<String> {
52        self.input.reset().push_str(text);
53        self.input.start_build()?;
54        self.source.rewrite(&mut self.input)?;
55        Ok(self.input.current().to_owned())
56    }
57}
58
59impl<D> TextNormalizer<D>
60where
61    D: DictionaryAccess,
62{
63    /// Create a text normalizer using the input-text plugins from a dictionary.
64    pub fn from_dictionary(dictionary: D) -> Self {
65        Self {
66            source: dictionary,
67            input: InputBuffer::new(),
68        }
69    }
70
71    pub fn normalize(&mut self, text: &str) -> SudachiResult<String> {
72        self.input.reset().push_str(text);
73        self.input.start_build()?;
74        rewrite_with_dictionary(&self.source, &mut self.input)?;
75        Ok(self.input.current().to_owned())
76    }
77}
78
79impl<D> TextNormalizer<Arc<D>>
80where
81    D: DictionaryAccess + ?Sized,
82{
83    /// Create a text normalizer using the input-text plugins from a shared dictionary handle.
84    pub fn from_shared_dictionary(dictionary: Arc<D>) -> Self {
85        TextNormalizer::from_dictionary(dictionary)
86    }
87}
88
89fn set_up_default_plugin(grammar: &Grammar) -> SudachiResult<DefaultInputTextPlugin> {
90    let mut plugin = DefaultInputTextPlugin::default();
91    let cfg = ConfigBuilder::empty().push_embedded().build();
92    plugin.set_up(
93        &serde_json::Value::Object(serde_json::Map::default()),
94        &cfg,
95        grammar,
96    )?;
97    Ok(plugin)
98}
99
100fn rewrite_with_dictionary<D>(dictionary: &D, input: &mut InputBuffer) -> SudachiResult<()>
101where
102    D: DictionaryAccess + ?Sized,
103{
104    for plugin in dictionary.input_text_plugins() {
105        plugin.rewrite(input)?;
106    }
107    Ok(())
108}
109
110#[cfg(test)]
111mod tests {
112    use serde_json::Value;
113
114    use super::*;
115    use crate::config::Config;
116    use crate::dic::lexicon_set::LexiconSet;
117    use crate::dic::LexiconAccess;
118    use crate::input_text::InputEditor;
119    use crate::plugin::oov::OovProviderPlugin;
120    use crate::plugin::path_rewrite::PathRewritePlugin;
121
122    struct ReplaceAllPlugin;
123
124    impl InputTextPlugin for ReplaceAllPlugin {
125        fn set_up(
126            &mut self,
127            _settings: &Value,
128            _config: &Config,
129            _grammar: &Grammar,
130        ) -> SudachiResult<()> {
131            Ok(())
132        }
133
134        #[allow(deprecated)]
135        fn rewrite_impl<'a>(
136            &'a self,
137            input: &InputBuffer,
138            mut edit: InputEditor<'a>,
139        ) -> SudachiResult<InputEditor<'a>> {
140            edit.replace_ref(0..input.current().len(), "rewritten");
141            Ok(edit)
142        }
143    }
144
145    struct MockDictionary {
146        grammar: Grammar<'static>,
147        input_text_plugins: Vec<Box<dyn InputTextPlugin + Sync + Send>>,
148        oov_provider_plugins: Vec<Box<dyn OovProviderPlugin + Sync + Send>>,
149        path_rewrite_plugins: Vec<Box<dyn PathRewritePlugin + Sync + Send>>,
150    }
151
152    impl MockDictionary {
153        fn new(input_text_plugins: Vec<Box<dyn InputTextPlugin + Sync + Send>>) -> Self {
154            Self {
155                grammar: Grammar::empty(),
156                input_text_plugins,
157                oov_provider_plugins: Vec::new(),
158                path_rewrite_plugins: Vec::new(),
159            }
160        }
161    }
162
163    impl LexiconAccess for MockDictionary {
164        fn lexicon(&self) -> &LexiconSet<'_> {
165            unimplemented!("text normalization does not use lexicon access")
166        }
167    }
168
169    impl DictionaryAccess for MockDictionary {
170        fn grammar(&self) -> &Grammar<'_> {
171            &self.grammar
172        }
173
174        fn input_text_plugins(&self) -> &[Box<dyn InputTextPlugin + Sync + Send>] {
175            &self.input_text_plugins
176        }
177
178        fn oov_provider_plugins(&self) -> &[Box<dyn OovProviderPlugin + Sync + Send>] {
179            &self.oov_provider_plugins
180        }
181
182        fn path_rewrite_plugins(&self) -> &[Box<dyn PathRewritePlugin + Sync + Send>] {
183            &self.path_rewrite_plugins
184        }
185    }
186
187    #[test]
188    fn default_normalizer_works() {
189        let mut normalizer = TextNormalizer::try_default().unwrap();
190
191        assert_eq!("abc", normalizer.normalize("ABC").unwrap());
192        assert_eq!("", normalizer.normalize("").unwrap());
193        assert_eq!("ガヴ", normalizer.normalize("ガウ゛").unwrap());
194    }
195
196    #[test]
197    fn dictionary_normalizer_uses_dictionary_plugins() {
198        let dictionary = MockDictionary::new(vec![Box::new(ReplaceAllPlugin)]);
199        let mut normalizer = TextNormalizer::from_dictionary(&dictionary);
200
201        assert_eq!("rewritten", normalizer.normalize("abc").unwrap());
202        assert_eq!("rewritten", normalizer.normalize("ABC").unwrap());
203    }
204
205    #[test]
206    fn shared_dictionary_normalizer_uses_dictionary_plugins() {
207        let dictionary = Arc::new(MockDictionary::new(vec![Box::new(ReplaceAllPlugin)]));
208        let mut normalizer = TextNormalizer::from_shared_dictionary(dictionary);
209
210        assert_eq!("rewritten", normalizer.normalize("abc").unwrap());
211        assert_eq!("rewritten", normalizer.normalize("ABC").unwrap());
212    }
213}