lindera_dictionary/builder/
user_dictionary.rs

1use std::collections::BTreeMap;
2use std::fs;
3use std::fs::File;
4use std::io;
5use std::io::Write;
6use std::path::Path;
7
8use byteorder::{LittleEndian, WriteBytesExt};
9use csv::StringRecord;
10use daachorse::DoubleArrayAhoCorasickBuilder;
11use log::debug;
12
13use crate::LinderaResult;
14use crate::dictionary::UserDictionary;
15use crate::dictionary::prefix_dictionary::{DaTrust, PrefixDictionary};
16use crate::error::LinderaErrorKind;
17use crate::viterbi::WordEntry;
18
19type StringRecordProcessor = Option<Box<dyn Fn(&StringRecord) -> LinderaResult<Vec<String>>>>;
20
21pub struct UserDictionaryBuilder {
22    user_dictionary_fields_num: usize,
23    dictionary_fields_num: usize,
24    default_word_cost: i16,
25    default_left_context_id: u16,
26    default_right_context_id: u16,
27    flexible_csv: bool,
28    user_dictionary_handler: StringRecordProcessor,
29}
30
31/// Options for [`UserDictionaryBuilder`]. Every field has a default, so
32/// [`Self::builder`] is infallible. Setters take `self` by value to keep the
33/// original owned-builder calling style.
34#[derive(Default)]
35pub struct UserDictionaryBuilderOptions {
36    user_dictionary_fields_num: Option<usize>,
37    dictionary_fields_num: Option<usize>,
38    default_word_cost: Option<i16>,
39    default_left_context_id: Option<u16>,
40    default_right_context_id: Option<u16>,
41    flexible_csv: Option<bool>,
42    user_dictionary_handler: StringRecordProcessor,
43}
44
45impl UserDictionaryBuilderOptions {
46    pub fn user_dictionary_fields_num(mut self, value: usize) -> Self {
47        self.user_dictionary_fields_num = Some(value);
48        self
49    }
50
51    pub fn dictionary_fields_num(mut self, value: usize) -> Self {
52        self.dictionary_fields_num = Some(value);
53        self
54    }
55
56    pub fn default_word_cost(mut self, value: i16) -> Self {
57        self.default_word_cost = Some(value);
58        self
59    }
60
61    pub fn default_left_context_id(mut self, value: u16) -> Self {
62        self.default_left_context_id = Some(value);
63        self
64    }
65
66    pub fn default_right_context_id(mut self, value: u16) -> Self {
67        self.default_right_context_id = Some(value);
68        self
69    }
70
71    pub fn flexible_csv(mut self, value: bool) -> Self {
72        self.flexible_csv = Some(value);
73        self
74    }
75
76    pub fn user_dictionary_handler(mut self, value: StringRecordProcessor) -> Self {
77        self.user_dictionary_handler = value;
78        self
79    }
80
81    pub fn builder(self) -> UserDictionaryBuilder {
82        UserDictionaryBuilder {
83            user_dictionary_fields_num: self.user_dictionary_fields_num.unwrap_or(3),
84            dictionary_fields_num: self.dictionary_fields_num.unwrap_or(12),
85            default_word_cost: self.default_word_cost.unwrap_or(-10000),
86            default_left_context_id: self.default_left_context_id.unwrap_or(0),
87            default_right_context_id: self.default_right_context_id.unwrap_or(0),
88            flexible_csv: self.flexible_csv.unwrap_or(true),
89            user_dictionary_handler: self.user_dictionary_handler,
90        }
91    }
92}
93
94impl UserDictionaryBuilder {
95    pub fn build(&self, input_file: &Path) -> LinderaResult<UserDictionary> {
96        debug!("reading {input_file:?}");
97
98        let mut rdr = csv::ReaderBuilder::new()
99            .has_headers(false)
100            .flexible(self.flexible_csv)
101            .from_path(input_file)
102            .map_err(|err| {
103                LinderaErrorKind::Io
104                    .with_error(anyhow::anyhow!(err))
105                    .add_context(format!(
106                        "Failed to open user dictionary CSV file: {input_file:?}"
107                    ))
108            })?;
109
110        let mut rows: Vec<StringRecord> = vec![];
111        for (line_num, result) in rdr.records().enumerate() {
112            let record = result.map_err(|err| {
113                LinderaErrorKind::Content
114                    .with_error(anyhow::anyhow!(err))
115                    .add_context(format!(
116                        "Failed to parse CSV record at line {} in file: {:?}",
117                        line_num + 1,
118                        input_file
119                    ))
120            })?;
121            rows.push(record);
122        }
123        // The cached variant computes the allocating surface key once per row
124        // instead of once per comparison.
125        rows.sort_by_cached_key(|row| row[0].to_string());
126
127        let mut word_entry_map: BTreeMap<String, Vec<WordEntry>> = BTreeMap::new();
128
129        for (row_id, row) in rows.iter().enumerate() {
130            let surface = row[0].to_string();
131            let word_cost = if row.len() == self.user_dictionary_fields_num {
132                self.default_word_cost
133            } else {
134                row[3].parse::<i16>().map_err(|_err| {
135                    LinderaErrorKind::Parse
136                        .with_error(anyhow::anyhow!("failed to parse word cost"))
137                        .add_context(format!(
138                            "Invalid word cost '{}' at row {} (surface: '{}')",
139                            &row[3],
140                            row_id + 1,
141                            &row[0]
142                        ))
143                })?
144            };
145            let (left_id, right_id) = if row.len() == self.user_dictionary_fields_num {
146                (self.default_left_context_id, self.default_right_context_id)
147            } else {
148                (
149                    row[1].parse::<u16>().map_err(|_err| {
150                        LinderaErrorKind::Parse
151                            .with_error(anyhow::anyhow!("failed to parse left context id"))
152                            .add_context(format!(
153                                "Invalid left context ID '{}' at row {} (surface: '{}')",
154                                &row[1],
155                                row_id + 1,
156                                &row[0]
157                            ))
158                    })?,
159                    row[2].parse::<u16>().map_err(|_err| {
160                        LinderaErrorKind::Parse
161                            .with_error(anyhow::anyhow!("failed to parse right context id"))
162                            .add_context(format!(
163                                "Invalid right context ID '{}' at row {} (surface: '{}')",
164                                &row[2],
165                                row_id + 1,
166                                &row[0]
167                            ))
168                    })?,
169                )
170            };
171
172            word_entry_map
173                .entry(surface)
174                .or_default()
175                .push(WordEntry::new(
176                    crate::viterbi::WordId::new(crate::viterbi::LexType::User, row_id as u32),
177                    word_cost,
178                    left_id,
179                    right_id,
180                ));
181        }
182
183        let mut words_data = Vec::<u8>::new();
184        let mut words_idx_data = Vec::<u8>::new();
185        for row in rows.iter() {
186            let word_detail = if row.len() == self.user_dictionary_fields_num {
187                if let Some(handler) = &self.user_dictionary_handler {
188                    handler(row)?
189                } else {
190                    row.iter()
191                        .skip(1)
192                        .map(|s| s.to_string())
193                        .collect::<Vec<String>>()
194                }
195            } else if row.len() >= self.dictionary_fields_num {
196                let mut tmp_word_detail = Vec::new();
197                for item in row.iter().skip(4) {
198                    tmp_word_detail.push(item.to_string());
199                }
200                tmp_word_detail
201            } else {
202                return Err(LinderaErrorKind::Content
203                    .with_error(anyhow::anyhow!(
204                        "user dictionary should be a CSV with {} or {}+ fields",
205                        self.user_dictionary_fields_num,
206                        self.dictionary_fields_num
207                    ))
208                    .add_context(format!(
209                        "Row {} has {} fields (surface: '{}')",
210                        rows.iter().position(|r| std::ptr::eq(r, row)).unwrap_or(0) + 1,
211                        row.len(),
212                        row.get(0).unwrap_or("<empty>")
213                    )));
214            };
215
216            let offset = words_data.len();
217            words_idx_data
218                .write_u32::<LittleEndian>(offset as u32)
219                .map_err(|err| {
220                    LinderaErrorKind::Io
221                        .with_error(anyhow::anyhow!(err))
222                        .add_context("Failed to write word offset to user dictionary words index")
223                })?;
224
225            // Store word details as null-separated string (like main dictionary)
226            let joined_details = word_detail.join("\0");
227            let joined_details_len = u32::try_from(joined_details.len()).map_err(|err| {
228                LinderaErrorKind::Serialize
229                    .with_error(anyhow::anyhow!(err))
230                    .add_context(format!(
231                        "Word details length too large: {} bytes for word '{}'",
232                        joined_details.len(),
233                        row.get(0).unwrap_or("<unknown>")
234                    ))
235            })?;
236
237            words_data
238                .write_u32::<LittleEndian>(joined_details_len)
239                .map_err(|err| {
240                    LinderaErrorKind::Serialize
241                        .with_error(anyhow::anyhow!(err))
242                        .add_context(
243                            "Failed to write word details length to user dictionary words data",
244                        )
245                })?;
246            words_data
247                .write_all(joined_details.as_bytes())
248                .map_err(|err| {
249                    LinderaErrorKind::Serialize
250                        .with_error(anyhow::anyhow!(err))
251                        .add_context("Failed to write word details to user dictionary words data")
252                })?;
253        }
254
255        let mut id = 0u32;
256
257        // building double array trie
258        let mut keyset: Vec<(&[u8], u32)> = vec![];
259        for (key, word_entries) in &word_entry_map {
260            let len = word_entries.len() as u32;
261            // 24bit for word ID, 8bit for variant count (up to 255 per surface),
262            // matching the system dictionary encoding. The legacy 5-bit user
263            // encoding (max 31 variants) was retired in v4.0.0; user dictionary
264            // `.bin` files built with v3 must be rebuilt from their CSV source.
265            let val = (id << 8) | len;
266            keyset.push((key.as_bytes(), val));
267            id += len;
268        }
269        let da_bytes = DoubleArrayAhoCorasickBuilder::new()
270            .build_with_values(keyset)
271            .map_err(|err| {
272                LinderaErrorKind::Build
273                    .with_error(anyhow::anyhow!(err))
274                    .add_context("Failed to build DoubleArray for user dictionary")
275            })?
276            .serialize();
277
278        // building values
279        let mut vals_data = Vec::<u8>::new();
280        for word_entries in word_entry_map.values() {
281            for word_entry in word_entries {
282                word_entry.serialize(&mut vals_data).map_err(|err| {
283                    LinderaErrorKind::Serialize
284                        .with_error(anyhow::anyhow!(err))
285                        .add_context(format!(
286                            "Failed to serialize user dictionary word entry (id: {})",
287                            word_entry.word_id().id()
288                        ))
289                })?;
290            }
291        }
292
293        let dict = PrefixDictionary::load(
294            da_bytes,
295            vals_data,
296            words_idx_data,
297            words_data,
298            false,
299            DaTrust::Untrusted,
300        )?;
301
302        Ok(UserDictionary { dict })
303    }
304}
305
306pub fn build_user_dictionary(user_dict: UserDictionary, output_file: &Path) -> LinderaResult<()> {
307    let parent_dir = match output_file.parent() {
308        Some(parent_dir) => parent_dir,
309        None => {
310            return Err(LinderaErrorKind::Io
311                .with_error(anyhow::anyhow!(
312                    "failed to get parent directory of output file"
313                ))
314                .add_context(format!("Invalid output file path: {output_file:?}")));
315        }
316    };
317    fs::create_dir_all(parent_dir).map_err(|err| {
318        LinderaErrorKind::Io
319            .with_error(anyhow::anyhow!(err))
320            .add_context(format!("Failed to create parent directory: {parent_dir:?}"))
321    })?;
322
323    let mut wtr = io::BufWriter::new(File::create(output_file).map_err(|err| {
324        LinderaErrorKind::Io
325            .with_error(anyhow::anyhow!(err))
326            .add_context(format!(
327                "Failed to create user dictionary output file: {output_file:?}"
328            ))
329    })?);
330    let bytes = rkyv::to_bytes::<rkyv::rancor::Error>(&user_dict).map_err(|err| {
331        LinderaErrorKind::Serialize
332            .with_error(anyhow::anyhow!(err))
333            .add_context(format!(
334                "Failed to serialize user dictionary to file: {output_file:?}"
335            ))
336    })?;
337    wtr.write_all(&bytes).map_err(|err| {
338        LinderaErrorKind::Io
339            .with_error(anyhow::anyhow!(err))
340            .add_context(format!(
341                "Failed to write user dictionary to file: {output_file:?}"
342            ))
343    })?;
344    wtr.flush().map_err(|err| {
345        LinderaErrorKind::Io
346            .with_error(anyhow::anyhow!(err))
347            .add_context(format!(
348                "Failed to flush user dictionary output file: {output_file:?}"
349            ))
350    })?;
351
352    Ok(())
353}
lindera_dictionary/builder/user_dictionary.rs

lindera_dictionary/builder/
user_dictionary.rs
