sudachi/analysis/
stateful_tokenizer.rs

1/*
2 *  Copyright (c) 2021-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 crate::analysis::created::CreatedWords;
18use crate::analysis::inner::{Node, NodeIdx};
19use crate::analysis::lattice::Lattice;
20use crate::analysis::node::{LatticeNode, ResultNode};
21use crate::analysis::stateless_tokenizer::{dump_path, split_path};
22use crate::analysis::Mode;
23use crate::dic::connect::ConnectionMatrix;
24use crate::dic::lexicon::LexiconEntry;
25use crate::dic::lexicon_set::LexiconSet;
26use crate::dic::subset::InfoSubset;
27use crate::dic::word_info::WordInfo;
28use crate::dic::DictionaryAccess;
29use crate::error::{SudachiError, SudachiResult};
30use crate::input_text::InputBuffer;
31use crate::plugin::oov::OovProviderPlugin;
32use crate::prelude::MorphemeList;
33
34pub struct StatefulTokenizer<D> {
35    dictionary: D,
36    input: InputBuffer,
37    debug: bool,
38    mode: Mode,
39    oov: Vec<Node>,
40    lattice: Lattice,
41    top_path_ids: Vec<NodeIdx>,
42    top_path: Option<Vec<ResultNode>>,
43    subset: InfoSubset,
44    /// Per-boundary dictionary-match cache, reused across sentences.
45    match_cache: Vec<Vec<LexiconEntry>>,
46    /// Use the pipelined + prefetched dictionary lookup in `build_lattice`.
47    pipelined_lookup: bool,
48}
49
50impl<D: DictionaryAccess + Clone> StatefulTokenizer<D> {
51    /// Get a clone of current dictionary
52    pub fn dict_clone(&self) -> D {
53        self.dictionary.clone()
54    }
55}
56
57impl<D: DictionaryAccess> StatefulTokenizer<D> {
58    /// Create a new non-debug stateful tokenizer
59    pub fn new(dic: D, mode: Mode) -> Self {
60        Self::create(dic, false, mode)
61    }
62
63    /// Create a new debug stateful tokenizer with the following options
64    pub fn create(dic: D, debug: bool, mode: Mode) -> Self {
65        Self {
66            dictionary: dic,
67            input: InputBuffer::default(),
68            debug,
69            mode,
70            oov: Vec::with_capacity(10),
71            lattice: Lattice::default(),
72            top_path_ids: Vec::new(),
73            top_path: Some(Vec::new()),
74            subset: InfoSubset::all(),
75            match_cache: Vec::new(),
76            pipelined_lookup: true,
77        }
78    }
79
80    /// Enable or disable the pipelined + prefetched dictionary lookup. Output is
81    /// identical to the scalar path; this only affects performance. Returns the
82    /// previous value.
83    pub fn set_pipelined_lookup(&mut self, enabled: bool) -> bool {
84        std::mem::replace(&mut self.pipelined_lookup, enabled)
85    }
86
87    /// Set debug flag and returns the current one
88    pub fn set_debug(&mut self, debug: bool) -> bool {
89        std::mem::replace(&mut self.debug, debug)
90    }
91
92    /// Set the analysis mode and returns the current one
93    pub fn set_mode(&mut self, mode: Mode) -> Mode {
94        self.subset |= match mode {
95            Mode::A => InfoSubset::SPLIT_A,
96            Mode::B => InfoSubset::SPLIT_B,
97            _ => InfoSubset::empty(),
98        };
99        std::mem::replace(&mut self.mode, mode)
100    }
101
102    /// Return current analysis mode
103    pub fn mode(&self) -> Mode {
104        self.mode
105    }
106
107    /// Analyzer will read only following [`WordInfo`] field subset
108    pub fn set_subset(&mut self, subset: InfoSubset) -> InfoSubset {
109        let mode_subset = match self.mode {
110            Mode::A => InfoSubset::SPLIT_A,
111            Mode::B => InfoSubset::SPLIT_B,
112            _ => InfoSubset::empty(),
113        };
114        let new_subset = (subset | mode_subset).normalize();
115        std::mem::replace(&mut self.subset, new_subset | mode_subset)
116    }
117
118    /// Prepare StatefulTokenizer for the next data.
119    /// Data must be written in the returned reference.
120    pub fn reset(&mut self) -> &mut String {
121        if let Some(p) = self.top_path.as_mut() {
122            p.clear()
123        }
124        self.oov.clear();
125        self.input.reset()
126    }
127
128    /// Borrow current dictionary
129    pub fn dict(&self) -> &D {
130        &self.dictionary
131    }
132
133    /// Perform the actual tokenization so the analysis result will be available
134    /// for consumption
135    pub fn do_tokenize(&mut self) -> SudachiResult<()> {
136        self.input.start_build()?;
137        self.rewrite_input()?;
138        self.input.build(self.dictionary.grammar())?;
139
140        if self.input.current().is_empty() {
141            return Ok(());
142        }
143
144        let debug = self.debug;
145
146        if debug {
147            println!("=== Input dump:\n{}", self.input.current());
148        }
149
150        self.build_lattice()?;
151
152        if debug {
153            println!("=== Lattice dump:");
154            let dict = &self.dictionary;
155            let mut writer = std::io::stdout();
156            self.lattice
157                .dump(&self.input, dict.grammar(), dict.lexicon(), &mut writer)?;
158        };
159
160        let mut path = self.resolve_best_path()?;
161
162        if debug {
163            println!("=== Before Rewriting:");
164            dump_path(&path);
165        };
166
167        for plugin in self.dictionary.path_rewrite_plugins() {
168            path = plugin.rewrite(&self.input, path, &self.lattice, self.dictionary.lexicon())?;
169        }
170
171        path = split_path(&self.dictionary, path, self.mode, self.subset, &self.input)?;
172
173        if debug {
174            println!("=== After Rewriting:");
175            dump_path(&path);
176            println!("===");
177        };
178
179        self.top_path = Some(path);
180
181        Ok(())
182    }
183
184    /// Resolve the path (as ResultNodes) with the smallest cost
185    fn resolve_best_path(&mut self) -> SudachiResult<Vec<ResultNode>> {
186        let lexset = self.dictionary.lexicon();
187        let mut path = self.top_path.take().unwrap_or_default();
188        self.lattice.fill_top_path(&mut self.top_path_ids);
189        self.top_path_ids.reverse();
190        for pid in self.top_path_ids.drain(..) {
191            let (inner, cost) = self.lattice.node(pid);
192            let wi = if inner.word_id().is_oov() {
193                let curr_slice = self.input.curr_slice_c(inner.char_range()).to_owned();
194                WordInfo::new_oov(
195                    inner.word_id().entry().as_raw() as u16,
196                    curr_slice.len() as i16,
197                    inner.word_id(),
198                    curr_slice,
199                )
200            } else {
201                lexset.get_word_info_subset(inner.word_id(), self.subset)?
202            };
203
204            let byte_begin = self.input.to_curr_byte_idx(inner.begin());
205            let byte_end = self.input.to_curr_byte_idx(inner.end());
206
207            path.push(ResultNode::new(
208                inner.clone(),
209                cost,
210                byte_begin as u16,
211                byte_end as u16,
212                wi,
213            ));
214        }
215        Ok(path)
216    }
217
218    /// Swap result data with the current analyzer
219    pub fn swap_result(
220        &mut self,
221        input: &mut InputBuffer,
222        result: &mut Vec<ResultNode>,
223        subset: &mut InfoSubset,
224    ) {
225        std::mem::swap(&mut self.input, input);
226        std::mem::swap(self.top_path.as_mut().unwrap(), result);
227        *subset = self.subset;
228    }
229
230    fn rewrite_input(&mut self) -> SudachiResult<()> {
231        for p in self.dictionary.input_text_plugins() {
232            p.rewrite(&mut self.input)?;
233        }
234        Ok(())
235    }
236
237    fn build_lattice(&mut self) -> SudachiResult<()> {
238        let mut builder = LatticeBuilder {
239            node_buffer: &mut self.oov,
240            lattice: &mut self.lattice,
241            matrix: self.dictionary.grammar().conn_matrix(),
242            oov_providers: self.dictionary.oov_provider_plugins(),
243            lexicon: self.dictionary.lexicon(),
244            input: &self.input,
245            match_cache: &mut self.match_cache,
246            pipelined: self.pipelined_lookup,
247        };
248        builder.build_lattice()
249    }
250
251    /// Consume the Tokenizer and produce MorphemeList
252    pub fn into_morpheme_list(self) -> SudachiResult<MorphemeList<D>> {
253        match self.top_path {
254            None => Err(SudachiError::EosBosDisconnect),
255            Some(path) => Ok(MorphemeList::from_components(
256                self.dictionary,
257                self.input,
258                path,
259                self.subset,
260            )),
261        }
262    }
263}
264
265// This structure is purely for Rust.
266// Otherwise splitting code into functions fails to compile with double borrow errors
267struct LatticeBuilder<'a> {
268    node_buffer: &'a mut Vec<Node>,
269    lattice: &'a mut Lattice,
270    matrix: &'a ConnectionMatrix<'a>,
271    input: &'a InputBuffer,
272    lexicon: &'a LexiconSet<'a>,
273    oov_providers: &'a [Box<dyn OovProviderPlugin + Sync + Send>],
274    match_cache: &'a mut Vec<Vec<LexiconEntry>>,
275    pipelined: bool,
276}
277
278impl<'a> LatticeBuilder<'a> {
279    #[inline]
280    fn build_lattice(&mut self) -> SudachiResult<()> {
281        self.lattice.reset(self.input.current_chars().len());
282        if self.pipelined {
283            self.build_lattice_pipelined()
284        } else {
285            self.build_lattice_scalar()
286        }
287    }
288
289    /// Original lattice builder: one scalar common-prefix walk per reachable
290    /// boundary, interleaved with node insertion.
291    #[inline]
292    fn build_lattice_scalar(&mut self) -> SudachiResult<()> {
293        let input_bytes = self.input.current().as_bytes();
294
295        for (ch_off, &byte_off) in self.input.curr_byte_offsets().iter().enumerate() {
296            if !self.lattice.has_previous_node(ch_off) {
297                continue;
298            }
299
300            self.node_buffer.clear();
301            let mut created = CreatedWords::default();
302            for e in self.lexicon.lookup(input_bytes, byte_off) {
303                // do we really need input.can_bow condition?
304                if (e.end < input_bytes.len()) && !self.input.can_bow(e.end) {
305                    continue;
306                }
307                let (left_id, right_id, cost) = self.lexicon.get_word_param(e.word_id);
308                let end_c = self.input.ch_idx(e.end);
309                let node = Node::new(
310                    ch_off as u16,
311                    end_c as u16,
312                    left_id as u16,
313                    right_id as u16,
314                    cost,
315                    e.word_id,
316                );
317                created = created.add_word((end_c - ch_off) as i64);
318                self.node_buffer.push(node.clone());
319                self.lattice.insert(node, self.matrix);
320            }
321
322            self.insert_oovs(ch_off, created)?;
323        }
324        self.lattice.connect_eos(self.matrix)?;
325
326        Ok(())
327    }
328
329    /// Pre-compute every boundary's dictionary matches with overlapped trie
330    /// memory latency, then insert nodes for reachable boundaries in scalar
331    /// order. The walks depend only on the input, so precomputing them all is
332    /// exact; reachability still gates insertion.
333    #[inline]
334    fn build_lattice_pipelined(&mut self) -> SudachiResult<()> {
335        let input_bytes = self.input.current().as_bytes();
336
337        {
338            let lexicon = self.lexicon;
339            let starts = self.input.curr_byte_offsets();
340            let cache = &mut *self.match_cache;
341            if cache.len() < starts.len() {
342                cache.resize_with(starts.len(), Vec::new);
343            }
344            for bucket in cache.iter_mut() {
345                bucket.clear();
346            }
347            lexicon.lookup_batch(input_bytes, starts, |bucket, entry| {
348                cache[bucket].push(entry);
349            });
350        }
351
352        let boundaries = self.input.curr_byte_offsets().len();
353        for ch_off in 0..boundaries {
354            if !self.lattice.has_previous_node(ch_off) {
355                continue;
356            }
357
358            self.node_buffer.clear();
359            let mut created = CreatedWords::default();
360            for entry in &self.match_cache[ch_off] {
361                let (word_id, end) = (entry.word_id, entry.end);
362                if (end < input_bytes.len()) && !self.input.can_bow(end) {
363                    continue;
364                }
365                let (left_id, right_id, cost) = self.lexicon.get_word_param(word_id);
366                let end_c = self.input.ch_idx(end);
367                let node = Node::new(
368                    ch_off as u16,
369                    end_c as u16,
370                    left_id as u16,
371                    right_id as u16,
372                    cost,
373                    word_id,
374                );
375                created = created.add_word((end_c - ch_off) as i64);
376                self.node_buffer.push(node.clone());
377                self.lattice.insert(node, self.matrix);
378            }
379
380            self.insert_oovs(ch_off, created)?;
381        }
382        self.lattice.connect_eos(self.matrix)?;
383
384        Ok(())
385    }
386
387    /// OOV handling shared by both lattice builders. Mirrors the original
388    /// in-loop logic exactly.
389    #[inline]
390    fn insert_oovs(&mut self, ch_off: usize, mut created: CreatedWords) -> SudachiResult<()> {
391        if self.input.can_oov_bow(ch_off) {
392            for provider in self.oov_providers {
393                created = self.provide_oovs(ch_off, created, provider.as_ref())?;
394            }
395        }
396
397        if created.is_empty() {
398            let provider = self.oov_providers.last().unwrap();
399            created = self.provide_oovs(ch_off, created, provider.as_ref())?;
400        }
401
402        if created.is_empty() {
403            return Err(SudachiError::EosBosDisconnect);
404        }
405        Ok(())
406    }
407
408    #[inline]
409    fn provide_oovs<P>(
410        &mut self,
411        char_offset: usize,
412        mut other: CreatedWords,
413        plugin: &P,
414    ) -> SudachiResult<CreatedWords>
415    where
416        P: OovProviderPlugin + 'a + ?Sized,
417    {
418        let start_size = self.node_buffer.len();
419        let num_provided = plugin.provide_oov(self.input, char_offset, other, self.node_buffer)?;
420        for idx in start_size..(start_size + num_provided) {
421            let node = self.node_buffer[idx].clone();
422            other = other.add_word(node.char_range().len() as i64);
423            self.lattice.insert(node, self.matrix);
424        }
425        Ok(other)
426    }
427}
sudachi/analysis/stateful_tokenizer.rs

sudachi/analysis/
stateful_tokenizer.rs
