sudachi/
text_normalizer.rs1use 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
27pub struct TextNormalizer<D = DefaultInputTextPlugin> {
32 source: D,
33 input: InputBuffer,
34}
35
36impl TextNormalizer<DefaultInputTextPlugin> {
37 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 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 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 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}