novel_segment/
segment.rs

1//! Segmenter.
2
3use crate::dict::{SetDict, SynonymDict, TableDict};
4use crate::error::{Error, Result};
5use crate::loader::{
6    load_dict_file, load_line_file, load_synonym_file, parse_dict_line, resolve_named,
7};
8use crate::optimizers::builtin_optimizers;
9use crate::options::{DoSegmentOptions, SegmentOptions};
10use crate::pipeline::{
11    default_optimizer_names, default_tokenizer_names, enabled, Optimizer, Tokenizer,
12};
13use crate::postag::POSTAG;
14use crate::text::{is_newline_only, split_sections};
15use crate::tokenizers::{
16    builtin_tokenizers, DEFAULT_MAX_CHUNK_COUNT, DEFAULT_MAX_CHUNK_COUNT_MIN,
17};
18use crate::word::{stringify_list, Word};
19use std::collections::HashMap;
20use std::path::{Path, PathBuf};
21use std::sync::{Mutex, OnceLock};
22
23/// Chinese word segmenter.
24pub struct Segment {
25    pub(crate) options: SegmentOptions,
26    pub(crate) table: TableDict,
27    pub(crate) wildcard: TableDict,
28    pub(crate) synonym: SynonymDict,
29    pub(crate) stopword: SetDict,
30    pub(crate) blacklist: SetDict,
31    pub(crate) blacklist_optimizer: SetDict,
32    pub(crate) blacklist_synonym: SetDict,
33    pub(crate) tokenizers: Vec<Box<dyn Tokenizer>>,
34    pub(crate) optimizers: Vec<Box<dyn Optimizer>>,
35    pub(crate) max_chunk_count: usize,
36    pub(crate) min_chunk_count: usize,
37    pub(crate) inited: bool,
38}
39
40impl Default for Segment {
41    fn default() -> Self {
42        Self::new(SegmentOptions::default())
43    }
44}
45
46impl Segment {
47    pub fn new(options: SegmentOptions) -> Self {
48        let max_chunk_count = options
49            .max_chunk_count
50            .filter(|n| *n > DEFAULT_MAX_CHUNK_COUNT_MIN)
51            .unwrap_or(DEFAULT_MAX_CHUNK_COUNT);
52        let min_chunk_count = options
53            .min_chunk_count
54            .filter(|n| *n > DEFAULT_MAX_CHUNK_COUNT_MIN)
55            .unwrap_or(DEFAULT_MAX_CHUNK_COUNT_MIN);
56        Self {
57            options,
58            table: TableDict::new(),
59            wildcard: TableDict::new(),
60            synonym: SynonymDict::default(),
61            stopword: SetDict::default(),
62            blacklist: SetDict::default(),
63            blacklist_optimizer: SetDict::default(),
64            blacklist_synonym: SetDict::default(),
65            tokenizers: Vec::new(),
66            optimizers: Vec::new(),
67            max_chunk_count,
68            min_chunk_count,
69            inited: false,
70        }
71    }
72
73    /// `new Segment(); useDefault()` — `auto_cjk` stays false.
74    pub fn with_default() -> Result<Self> {
75        let mut seg = Self::new(SegmentOptions::default());
76        seg.use_default()?;
77        Ok(seg)
78    }
79
80    /// Test/CLI profile: `autoCjk`, `all_mod`, `convertSynonym`.
81    pub fn with_novel_default() -> Result<Self> {
82        let mut seg = Self::new(SegmentOptions {
83            auto_cjk: true,
84            all_mod: true,
85            options_do_segment: DoSegmentOptions::convert_synonym(),
86            ..Default::default()
87        });
88        seg.use_default()?;
89        Ok(seg)
90    }
91
92    /// `word.novel.test.ts`: createSegment then wipe SYNONYM and reload with `nodeNovelMode`.
93    pub fn with_node_novel_default() -> Result<Self> {
94        let mut seg = Self::new(SegmentOptions {
95            auto_cjk: true,
96            all_mod: true,
97            node_novel_mode: true,
98            options_do_segment: DoSegmentOptions::convert_synonym(),
99            ..Default::default()
100        });
101        seg.use_default()?;
102        seg.clear_synonym_dict();
103        seg.use_default_synonym_dict(true)?;
104        Ok(seg)
105    }
106
107    pub fn use_default(&mut self) -> Result<&mut Self> {
108        if !self.options.nomod {
109            self.use_default_modules();
110        }
111        if !self.options.nodict {
112            self.use_default_dicts()?;
113        }
114        self.inited = true;
115        Ok(self)
116    }
117
118    pub fn use_default_modules(&mut self) -> &mut Self {
119        let tok_names = default_tokenizer_names(self.options.all_mod);
120        let opt_names = default_optimizer_names(self.options.all_mod);
121        let disabled = &self.options.disable_modules;
122        self.tokenizers = builtin_tokenizers()
123            .into_iter()
124            .filter(|t| tok_names.contains(&t.name()) && enabled(t.name(), disabled))
125            .collect();
126        self.optimizers = builtin_optimizers()
127            .into_iter()
128            .filter(|t| opt_names.contains(&t.name()) && enabled(t.name(), disabled))
129            .collect();
130        self
131    }
132
133    pub fn use_tokenizer(&mut self, tok: Box<dyn Tokenizer>) -> &mut Self {
134        self.tokenizers.push(tok);
135        self.inited = true;
136        self
137    }
138
139    pub fn use_optimizer(&mut self, opt: Box<dyn Optimizer>) -> &mut Self {
140        self.optimizers.push(opt);
141        self.inited = true;
142        self
143    }
144
145    #[cfg(feature = "default-dict")]
146    fn segment_root(&self) -> PathBuf {
147        novel_segment_dict::segment_dict_root()
148    }
149
150    #[cfg(not(feature = "default-dict"))]
151    fn segment_root(&self) -> PathBuf {
152        PathBuf::from("dict/segment")
153    }
154
155    pub fn use_default_dicts(&mut self) -> Result<&mut Self> {
156        #[cfg(feature = "default-dict")]
157        {
158            if let Some(cached) = lookup_cache(&self.options) {
159                self.apply_cached_dicts(&cached);
160                return Ok(self);
161            }
162        }
163        self.load_dict("char")?;
164        self.load_dict("pangu/phrases")?;
165        self.load_dict("pangu/phrases2")?;
166        self.load_dict("phrases/*")?;
167        self.load_dict("dict")?;
168        self.load_dict("dict2")?;
169        self.load_dict("dict3")?;
170        self.load_dict("dict4")?;
171        self.load_dict("pangu/dict005")?;
172        self.load_dict("pangu/dict006")?;
173        self.load_dict("dict_synonym/*")?;
174        self.load_stopword_dict("stopword")?;
175        self.load_dict("lazy/dict_synonym")?;
176        self.load_dict("names/*")?;
177        self.load_dict("lazy/*")?;
178        self.load_dict("pangu/num")?;
179        self.load_dict("lazy/badword")?;
180        self.load_dict_as("pangu/wildcard", DictKind::Wildcard, true)?;
181        self.load_synonym_dict("synonym", true)?;
182        self.load_synonym_dict("zht.synonym", false)?;
183        if self.options.node_novel_mode {
184            self.load_synonym_dict("badword.synonym", false)?;
185            self.load_synonym_dict("zht.common.synonym", false)?;
186        }
187        self.load_blacklist_dict("blacklist")?;
188        self.load_blacklist_optimizer_dict("blacklist.name")?;
189        self.load_blacklist_synonym_dict("blacklist.synonym")?;
190        self.do_blacklist();
191        #[cfg(feature = "default-dict")]
192        {
193            store_cache(self);
194        }
195        Ok(self)
196    }
197
198    pub fn load_dict(&mut self, name: &str) -> Result<&mut Self> {
199        self.load_dict_as(name, DictKind::Table, false)
200    }
201
202    pub fn load_dict_file(&mut self, path: impl AsRef<Path>) -> Result<&mut Self> {
203        let path = path.as_ref();
204        let rows = load_dict_file(path)?;
205        for row in rows {
206            self.table.add_with_cjk(&row.w, row.p, row.f, self.options.auto_cjk);
207        }
208        self.inited = true;
209        Ok(self)
210    }
211
212    fn load_dict_as(&mut self, name: &str, kind: DictKind, lower: bool) -> Result<&mut Self> {
213        let files = resolve_named(&self.segment_root(), name)?;
214        for file in files {
215            let rows = load_dict_file(&file)?;
216            for mut row in rows {
217                if lower {
218                    row.w = row.w.to_lowercase();
219                }
220                match kind {
221                    DictKind::Table => {
222                        self.table.add_with_cjk(&row.w, row.p, row.f, self.options.auto_cjk);
223                    }
224                    DictKind::Wildcard => {
225                        self.wildcard.add(row.w, row.p, row.f, true);
226                    }
227                }
228            }
229        }
230        self.inited = true;
231        Ok(self)
232    }
233
234    pub fn load_synonym_dict(&mut self, name: &str, skip_exists: bool) -> Result<&mut Self> {
235        #[cfg(feature = "default-dict")]
236        let root = novel_segment_dict::synonym_dict_root();
237        #[cfg(not(feature = "default-dict"))]
238        let root = PathBuf::from("dict/synonym");
239        let files = match resolve_named(&root, name) {
240            Ok(f) => f,
241            Err(_) => resolve_named(&self.segment_root(), name)?,
242        };
243        for file in files {
244            for (canon, vars) in load_synonym_file(&file)? {
245                self.synonym.add(&canon, &vars, skip_exists);
246                // Phrase synonyms must be tokenizable (word.novel: 恐怖分子 → 恐怖份子).
247                if let Some(e) = self.table.table.get(&canon).cloned() {
248                    for v in &vars {
249                        if !self.table.table.contains_key(v) {
250                            self.table.add(v, e.p, e.f, false);
251                        }
252                    }
253                }
254            }
255        }
256        self.inited = true;
257        Ok(self)
258    }
259
260    pub fn clear_synonym_dict(&mut self) -> &mut Self {
261        self.synonym = SynonymDict::default();
262        self
263    }
264
265    /// JS `useDefaultSynonymDict`.
266    pub fn use_default_synonym_dict(&mut self, node_novel_mode: bool) -> Result<&mut Self> {
267        self.load_synonym_dict("synonym", true)?;
268        self.load_synonym_dict("zht.synonym", false)?;
269        if node_novel_mode {
270            self.load_synonym_dict("badword.synonym", false)?;
271            self.load_synonym_dict("zht.common.synonym", false)?;
272        }
273        Ok(self)
274    }
275
276    pub fn load_stopword_dict(&mut self, name: &str) -> Result<&mut Self> {
277        #[cfg(feature = "default-dict")]
278        let root = novel_segment_dict::stopword_dict_root();
279        #[cfg(not(feature = "default-dict"))]
280        let root = PathBuf::from("dict/stopword");
281        let files = match resolve_named(&root, name) {
282            Ok(f) => f,
283            Err(_) => resolve_named(&self.segment_root(), name)?,
284        };
285        for file in files {
286            for line in load_line_file(&file)? {
287                self.stopword.add(line);
288            }
289        }
290        Ok(self)
291    }
292
293    pub fn load_blacklist_dict(&mut self, name: &str) -> Result<&mut Self> {
294        self.load_set_dict(name, BlackKind::Main)
295    }
296
297    pub fn load_blacklist_optimizer_dict(&mut self, name: &str) -> Result<&mut Self> {
298        self.load_set_dict(name, BlackKind::Optimizer)
299    }
300
301    pub fn load_blacklist_synonym_dict(&mut self, name: &str) -> Result<&mut Self> {
302        self.load_set_dict(name, BlackKind::Synonym)
303    }
304
305    fn load_set_dict(&mut self, name: &str, kind: BlackKind) -> Result<&mut Self> {
306        #[cfg(feature = "default-dict")]
307        let root = novel_segment_dict::blacklist_dict_root();
308        #[cfg(not(feature = "default-dict"))]
309        let root = PathBuf::from("dict/blacklist");
310        let files = match resolve_named(&root, name) {
311            Ok(f) => f,
312            Err(_) => resolve_named(&self.segment_root(), name)?,
313        };
314        for file in files {
315            for line in load_line_file(&file)? {
316                match kind {
317                    BlackKind::Main => self.blacklist.add(line),
318                    BlackKind::Optimizer => self.blacklist_optimizer.add(line),
319                    BlackKind::Synonym => self.blacklist_synonym.add(line),
320                }
321            }
322        }
323        Ok(self)
324    }
325
326    pub fn do_blacklist(&mut self) -> &mut Self {
327        let keys: Vec<String> = self.blacklist.table.iter().cloned().collect();
328        for k in keys {
329            self.table.remove(&k);
330        }
331        self
332    }
333
334    /// Remove a word from TABLE (JS `addBlacklist`).
335    pub fn add_blacklist(&mut self, word: &str) -> &mut Self {
336        if !word.is_empty() {
337            self.blacklist.add(word.to_string());
338            self.table.remove(word);
339        }
340        self
341    }
342
343    /// Add a dictionary entry (`詞|詞性|詞權值` or just the word).
344    pub fn add_word(&mut self, spec: &str, p: Option<u32>, f: Option<f64>) -> Result<&mut Self> {
345        if let Some(row) = parse_dict_line(spec) {
346            self.table
347                .add_with_cjk(&row.w, row.p, row.f, self.options.auto_cjk);
348            return Ok(self);
349        }
350        if spec.trim().is_empty() {
351            return Err(Error::InvalidInput(spec.to_string()));
352        }
353        self.table.add_with_cjk(
354            spec.trim(),
355            p.unwrap_or(0),
356            f.unwrap_or(0.0),
357            self.options.auto_cjk,
358        );
359        Ok(self)
360    }
361
362    pub fn add_synonym(&mut self, canonical: &str, variants: &[&str]) -> &mut Self {
363        let vars: Vec<String> = variants.iter().map(|s| (*s).to_string()).collect();
364        self.synonym.add(canonical, &vars, false);
365        if let Some(e) = self.table.table.get(canonical).cloned() {
366            for v in &vars {
367                if !self.table.table.contains_key(v) {
368                    self.table.add(v, e.p, e.f, false);
369                }
370            }
371        }
372        self
373    }
374
375    pub fn do_segment(&self, text: &str, options: DoSegmentOptions) -> Vec<Word> {
376        let options = self.options.options_do_segment.merge(&options);
377        let disabled = if options.disable_modules.is_empty() {
378            &self.options.disable_modules
379        } else {
380            &options.disable_modules
381        };
382        let toks: Vec<&dyn Tokenizer> = self
383            .tokenizers
384            .iter()
385            .filter(|t| enabled(t.name(), disabled))
386            .map(|t| t.as_ref())
387            .collect();
388        let opts: Vec<&dyn Optimizer> = self
389            .optimizers
390            .iter()
391            .filter(|t| enabled(t.name(), disabled))
392            .map(|t| t.as_ref())
393            .collect();
394
395        let mut ret = Vec::new();
396        for section in split_sections(text) {
397            if is_newline_only(&section) {
398                ret.push(Word::new(section));
399                continue;
400            }
401            if section.is_empty() {
402                continue;
403            }
404            let mut words = vec![Word::new(section)];
405            for t in &toks {
406                words = t.split(words, self);
407            }
408            for o in &opts {
409                words = o.do_optimize(words, self);
410            }
411            ret.extend(words);
412        }
413
414        if options.strip_punctuation_flag() {
415            ret.retain(|w| w.pos() != POSTAG::D_W);
416        }
417        if options.convert_synonym_flag() {
418            ret = convert_synonym(ret, self);
419        }
420        if options.strip_stopword_flag() {
421            ret.retain(|w| !self.stopword.contains(&w.w));
422        }
423        if options.strip_space_flag() {
424            ret.retain(|w| !w.w.chars().all(char::is_whitespace));
425        }
426        ret
427    }
428
429    pub fn do_segment_simple(&self, text: &str, mut options: DoSegmentOptions) -> Vec<String> {
430        options.simple = Some(true);
431        stringify_list(&self.do_segment(text, options))
432    }
433
434    fn apply_cached_dicts(&mut self, cached: &CachedDicts) {
435        self.table = cached.table.clone();
436        self.wildcard = cached.wildcard.clone();
437        self.synonym = cached.synonym.clone();
438        self.stopword = cached.stopword.clone();
439        self.blacklist = cached.blacklist.clone();
440        self.blacklist_optimizer = cached.blacklist_optimizer.clone();
441        self.blacklist_synonym = cached.blacklist_synonym.clone();
442        self.inited = true;
443    }
444}
445
446#[derive(Clone, Copy)]
447enum DictKind {
448    Table,
449    Wildcard,
450}
451
452enum BlackKind {
453    Main,
454    Optimizer,
455    Synonym,
456}
457
458fn convert_synonym(mut words: Vec<Word>, seg: &Segment) -> Vec<Word> {
459    loop {
460        let mut count = 0;
461        let mut next = Vec::with_capacity(words.len());
462        for item in words {
463            let w = item.w.clone();
464            let mut nw: Option<String> = None;
465            if let Some(s) = seg.synonym.get(&w) {
466                nw = Some(s.to_string());
467            } else if item.auto_create && item.ow.is_none() {
468                if let Some(parts) = &item.m {
469                    if !parts.is_empty() {
470                        let mut joined = String::new();
471                        let mut hit = false;
472                        for b in parts {
473                            if let Some(s) = seg.synonym.get(&b.w) {
474                                joined.push_str(s);
475                                hit = true;
476                            } else {
477                                joined.push_str(&b.w);
478                            }
479                        }
480                        if hit {
481                            nw = Some(joined);
482                        }
483                    }
484                }
485            }
486            if let Some(nw) = nw {
487                count += 1;
488                let mut p = item.pos();
489                if let Some(e) = seg.table.table.get(&w) {
490                    if e.p != 0 {
491                        p = e.p;
492                    }
493                }
494                if p & POSTAG::BAD != 0 {
495                    p ^= POSTAG::BAD;
496                }
497                next.push(Word {
498                    ow: Some(w),
499                    op: item.p,
500                    w: nw,
501                    p: Some(p),
502                    ..item
503                });
504            } else {
505                next.push(item);
506            }
507        }
508        words = next;
509        if count == 0 {
510            break;
511        }
512    }
513    words
514}
515
516#[derive(Clone)]
517struct CachedDicts {
518    table: TableDict,
519    wildcard: TableDict,
520    synonym: SynonymDict,
521    stopword: SetDict,
522    blacklist: SetDict,
523    blacklist_optimizer: SetDict,
524    blacklist_synonym: SetDict,
525}
526
527#[derive(Clone, PartialEq, Eq, Hash)]
528struct CacheKey {
529    auto_cjk: bool,
530    node_novel_mode: bool,
531    dict_root: String,
532}
533
534static DICT_CACHES: OnceLock<Mutex<HashMap<CacheKey, CachedDicts>>> = OnceLock::new();
535
536fn cache_key(options: &SegmentOptions) -> CacheKey {
537    let dict_root = {
538        #[cfg(feature = "default-dict")]
539        {
540            novel_segment_dict::dict_root().display().to_string()
541        }
542        #[cfg(not(feature = "default-dict"))]
543        {
544            String::from("dict")
545        }
546    };
547    CacheKey {
548        auto_cjk: options.auto_cjk,
549        node_novel_mode: options.node_novel_mode,
550        dict_root,
551    }
552}
553
554#[cfg(feature = "default-dict")]
555fn lookup_cache(options: &SegmentOptions) -> Option<CachedDicts> {
556    let map = DICT_CACHES.get_or_init(|| Mutex::new(HashMap::new()));
557    let guard = map.lock().ok()?;
558    guard.get(&cache_key(options)).cloned()
559}
560
561#[cfg(feature = "default-dict")]
562fn store_cache(seg: &Segment) {
563    let map = DICT_CACHES.get_or_init(|| Mutex::new(HashMap::new()));
564    if let Ok(mut guard) = map.lock() {
565        guard.insert(
566            cache_key(&seg.options),
567            CachedDicts {
568                table: seg.table.clone(),
569                wildcard: seg.wildcard.clone(),
570                synonym: seg.synonym.clone(),
571                stopword: seg.stopword.clone(),
572                blacklist: seg.blacklist.clone(),
573                blacklist_optimizer: seg.blacklist_optimizer.clone(),
574                blacklist_synonym: seg.blacklist_synonym.clone(),
575            },
576        );
577    }
578}
579
580#[cfg(test)]
581mod tests {
582    use super::*;
583
584    #[test]
585    fn convert_compound_parts() {
586        let mut seg = Segment::new(SegmentOptions::default());
587        seg.add_synonym("標準", &["錯字"]);
588        let mut item = Word::new("錯字甲");
589        item.auto_create = true;
590        item.m = Some(vec![Word::new("錯字"), Word::new("甲")]);
591        let out = convert_synonym(vec![item], &seg);
592        assert_eq!(out[0].w, "標準甲");
593        assert_eq!(out[0].ow.as_deref(), Some("錯字甲"));
594    }
595}
novel_segment/segment.rs

novel_segment/
segment.rs
