Skip to main content

lindera_trainer/
model.rs

1use std::collections::BTreeMap;
2use std::io::{Read, Write};
3use std::num::NonZeroU32;
4
5use anyhow::Result;
6use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};
7use serde::{Deserialize, Serialize};
8
9use crate::corpus::Word;
10use lindera_dictionary::viterbi::{LexType, WordEntry, WordId};
11
12/// Convert CRF weight to MeCab-compatible cost using cost-factor.
13///
14/// MeCab's `tocost(d, n)` formula: `clamp(-n * d, -32767, 32767)`
15/// where `d` is the cost-factor and `n` is the CRF weight.
16fn tocost(weight: f64, cost_factor: i32) -> i16 {
17    let raw = -(cost_factor as f64) * weight;
18    raw.round().clamp(i16::MIN as f64, i16::MAX as f64) as i16
19}
20
21/// Returns whether `bytes` begins a JSON object, ignoring leading whitespace.
22///
23/// Used only to decide whether a payload that rkyv already rejected is worth
24/// handing to the legacy JSON reader. rkyv is always tried first, so a false
25/// positive here cannot misroute a valid model.
26///
27/// # 引数
28///
29/// * `bytes` - Raw model file contents.
30///
31/// # 戻り値
32///
33/// `true` when the first non-whitespace byte is `{`.
34fn looks_like_json_object(bytes: &[u8]) -> bool {
35    bytes
36        .iter()
37        .find(|byte| !byte.is_ascii_whitespace())
38        .is_some_and(|byte| *byte == b'{')
39}
40
41/// Calculate the optimal cost factor from actual model weights.
42///
43/// This ensures the full i16 range is utilized, preserving relative differences
44/// between weights. The cost factor is computed as `i16::MAX / max_abs_weight`
45/// so that the largest weight maps to the boundary of the i16 range.
46fn calculate_cost_factor(merged_model: &lindera_crf::MergedModel) -> i32 {
47    let mut weight_abs_max = 0f64;
48
49    // Find maximum absolute weight from unigram feature sets
50    for feature_set in &merged_model.feature_sets {
51        weight_abs_max = weight_abs_max.max(feature_set.weight.abs());
52    }
53
54    // Find maximum absolute weight from connection cost matrix
55    for hm in &merged_model.matrix {
56        for &w in hm.values() {
57            weight_abs_max = weight_abs_max.max(w.abs());
58        }
59    }
60
61    if weight_abs_max > f64::EPSILON {
62        (f64::from(i16::MAX) / weight_abs_max) as i32
63    } else {
64        700 // MeCab default fallback
65    }
66}
67
68/// Feature set information extracted from CRF training
69#[derive(Serialize, Deserialize, Archive, RkyvSerialize, RkyvDeserialize, Clone, Debug)]
70
71pub struct FeatureSetInfo {
72    /// Left connection ID learned from CRF training
73    pub left_id: u32,
74    /// Right connection ID learned from CRF training
75    pub right_id: u32,
76    /// Feature weight learned from CRF training
77    pub weight: f64,
78}
79
80/// Trained model with weights and configuration.
81#[derive(Serialize, Deserialize, Archive, RkyvSerialize, RkyvDeserialize)]
82
83pub struct SerializableModel {
84    /// Feature weights from CRF training
85    pub feature_weights: Vec<f64>,
86    /// Label information
87    pub labels: Vec<String>,
88    /// Part-of-speech information for each label
89    pub pos_info: Vec<String>,
90    /// Model metadata
91    pub metadata: ModelMetadata,
92    /// Connection cost matrix: (right_id, left_id) -> cost.
93    ///
94    /// Ordered maps rather than hashed ones, at both levels. rkyv archives a
95    /// `HashMap` as a SwissTable whose slot assignment resolves collisions in
96    /// source iteration order, and `HashMap` iteration order is seeded per
97    /// instance -- so serializing the same trained model twice produced files
98    /// of identical length with scattered byte differences (#974). A
99    /// `BTreeMap` serializes from its key-ordered iterator into a tree whose
100    /// shape depends only on the entry count, which makes `model.dat`
101    /// byte-reproducible. No stored value changes: both levels are only ever
102    /// filled key-by-key from `Vec`-ordered sources.
103    pub connection_matrix: BTreeMap<usize, BTreeMap<usize, f64>>,
104    /// Maximum left connection ID
105    pub max_left_id: usize,
106    /// Maximum right connection ID
107    pub max_right_id: usize,
108    /// Feature set information (left_id, right_id, weight) for each label
109    pub feature_sets: Vec<FeatureSetInfo>,
110    /// Unknown word category names (from char.def)
111    pub unk_category_names: Vec<String>,
112    /// Unknown word category features (from unk.def).
113    ///
114    /// Ordered for the same reason as
115    /// [`SerializableModel::connection_matrix`]: rkyv's archived `HashMap`
116    /// layout follows source iteration order (#974).
117    pub unk_categories: BTreeMap<String, String>,
118    /// Raw content of the character definition file (char.def)
119    pub char_def_content: String,
120    /// Raw content of the feature definition file (feature.def)
121    pub feature_def_content: String,
122    /// Raw content of the rewrite rule definition file (rewrite.def)
123    pub rewrite_def_content: String,
124    /// Cost factor for weight-to-cost conversion (MeCab default: 700)
125    pub cost_factor: i32,
126    /// Left context ID to feature string mapping (for left-id.def)
127    pub left_id_map: Vec<(u32, String)>,
128    /// Right context ID to feature string mapping (for right-id.def)
129    pub right_id_map: Vec<(u32, String)>,
130}
131
132#[derive(Serialize, Deserialize, Archive, RkyvSerialize, RkyvDeserialize)]
133
134pub struct ModelMetadata {
135    pub version: String,
136    pub regularization: f64,
137    pub iterations: u64,
138    pub feature_count: usize,
139    pub label_count: usize,
140}
141
142/// Trained model.
143pub struct Model {
144    pub(crate) raw_model: lindera_crf::RawModel,
145    pub(crate) config: super::config::TrainerConfig,
146    pub(crate) feature_weights: Vec<f64>,
147    pub(crate) labels: Vec<String>,
148    pub(crate) user_entries: Vec<(Word, WordEntry, NonZeroU32)>,
149    pub(crate) merged_model: Option<lindera_crf::MergedModel>,
150    pub(crate) regularization_cost: f64,
151    pub(crate) max_iterations: u64,
152}
153
154impl Model {
155    /// Creates a new model with metadata.
156    pub(crate) fn new_with_metadata(
157        raw_model: lindera_crf::RawModel,
158        config: super::config::TrainerConfig,
159        feature_weights: Vec<f64>,
160        labels: Vec<String>,
161        regularization_cost: f64,
162        max_iterations: u64,
163    ) -> Self {
164        Self {
165            raw_model,
166            config,
167            feature_weights,
168            labels,
169            user_entries: Vec::new(),
170            merged_model: None,
171            regularization_cost,
172            max_iterations,
173        }
174    }
175
176    /// Reads the user-defined lexicon file.
177    ///
178    /// If you want to assign parameters to the user-defined lexicon file, you need to call this
179    /// function before exporting the dictionary. The model overwrites the parameter only when it
180    /// is `0,0,0`. Otherwise, the parameter is used as is.
181    ///
182    /// # Arguments
183    ///
184    /// * `rdr` - Read sink of the user-defined lexicon file.
185    pub fn read_user_lexicon<R: Read>(&mut self, rdr: R) -> Result<()> {
186        use std::io::BufRead;
187        use std::io::BufReader;
188
189        self.merged_model = None;
190        let reader = BufReader::new(rdr);
191
192        for line in reader.lines() {
193            let line = line?;
194            if line.trim().is_empty() || line.starts_with('#') {
195                continue;
196            }
197
198            let parts: Vec<&str> = line.split(',').collect();
199            if parts.len() >= 5 {
200                let surface = parts[0];
201                let left_id = parts[1].parse::<i32>().unwrap_or(0);
202                let right_id = parts[2].parse::<i32>().unwrap_or(0);
203                let cost = parts[3].parse::<i16>().unwrap_or(0);
204                let features = parts[4..].join(",");
205
206                let word = Word::new(surface, &features);
207
208                // Create word ID for user dictionary entry
209                let word_id = WordId::new(LexType::User, self.user_entries.len() as u32);
210
211                let entry = WordEntry::new(word_id, cost, left_id as u16, right_id as u16);
212
213                // Extract features and create feature set for this user entry
214                let first_char = surface.chars().next().unwrap_or('\0');
215                let cate_id = self.get_category_id(first_char);
216
217                // Apply dictionary rewriter to get ufeature, lfeature, rfeature
218                let (ufeature, lfeature, rfeature) =
219                    self.config.dictionary_rewriter.rewrite(&features);
220                let u_vec: Vec<&str> = ufeature.split(',').collect();
221                let l_vec: Vec<&str> = lfeature.split(',').collect();
222                let r_vec: Vec<&str> = rfeature.split(',').collect();
223
224                let unigram_features = self
225                    .config
226                    .feature_extractor
227                    .extract_unigram_feature_ids(&u_vec, cate_id);
228                let left_features = self
229                    .config
230                    .feature_extractor
231                    .extract_left_feature_ids(&l_vec);
232                let right_features = self
233                    .config
234                    .feature_extractor
235                    .extract_right_feature_ids(&r_vec);
236
237                let _feature_set = lindera_crf::FeatureSet::new(
238                    &unigram_features,
239                    &right_features,
240                    &left_features,
241                );
242                // TODO: Integrate feature_set into provider for proper user lexicon feature handling
243                // Currently, we cannot access the provider from this context, which limits
244                // the integration of user lexicon features into the trained model.
245                // This should be refactored to allow proper feature integration.
246
247                // Create a label ID without modifying the provider
248                // Since we can't clone the provider, we'll use a fixed ID based on entry count
249                let label_id = NonZeroU32::new(1000000 + self.user_entries.len() as u32 + 1)
250                    .ok_or_else(|| anyhow::anyhow!("Failed to create label ID"))?;
251
252                self.user_entries.push((word, entry, label_id));
253            }
254        }
255
256        Ok(())
257    }
258
259    fn get_category_id(&self, ch: char) -> u32 {
260        // Use CharacterDefinition to map character to category ID
261        // This works for any dictionary (IPADIC, UniDic, ko-dic, CC-CEDICT, etc.)
262        let char_def = &self.config.dict.character_definition;
263        let categories = char_def.lookup_categories(ch);
264
265        // Return the first category ID, or 0 (DEFAULT) if no categories match
266        if !categories.is_empty() {
267            categories[0].0 as u32
268        } else {
269            0 // DEFAULT category
270        }
271    }
272
273    /// Writes the model to a writer.
274    pub fn write_model<W: Write>(&self, writer: &mut W) -> Result<()> {
275        let feature_weights = self.feature_weights.clone();
276
277        // Extract connection cost matrix from the trained model
278        let merged_model = self.raw_model.merge()?;
279
280        // Compute optimal cost factor from actual weight range
281        let cost_factor = calculate_cost_factor(&merged_model);
282
283        let mut connection_matrix: BTreeMap<usize, BTreeMap<usize, f64>> = BTreeMap::new();
284        let mut max_left_id = 0;
285        let mut max_right_id = 0;
286
287        for (right_id, left_map) in merged_model.matrix.iter().enumerate() {
288            max_right_id = max_right_id.max(right_id);
289            let mut inner_map = BTreeMap::new();
290
291            for (&left_id, &weight) in left_map.iter() {
292                max_left_id = max_left_id.max(left_id as usize);
293                inner_map.insert(left_id as usize, weight);
294            }
295
296            if !inner_map.is_empty() {
297                connection_matrix.insert(right_id, inner_map);
298            }
299        }
300
301        // Extract feature_sets information from merged_model
302        let feature_sets: Vec<FeatureSetInfo> = merged_model
303            .feature_sets
304            .iter()
305            .map(|fs| FeatureSetInfo {
306                left_id: fs.left_id.get(),
307                right_id: fs.right_id.get(),
308                weight: fs.weight,
309            })
310            .collect();
311
312        // Extract unknown word category information
313        let char_def = &self.config.dict.character_definition;
314        let unk_category_names: Vec<String> = char_def
315            .categories()
316            .iter()
317            .map(|s| s.to_string())
318            .collect();
319
320        // Build left/right ID maps from merged model's connection IDs
321        // before feature_sets is moved into SerializableModel
322        // Build left/right ID maps from merged model's connection IDs.
323        // Use feature strings (POS info) from config.features, not surface forms.
324        let features = &self.config.features;
325        let unk_categories = &self.config.unk_categories;
326        let unk_start_idx = self
327            .labels
328            .len()
329            .saturating_sub(self.config.dict.character_definition.categories().len());
330
331        let get_feature_string = |i: usize| -> String {
332            if i < unk_start_idx {
333                // Dictionary entry: use feature string from config.features
334                if i < features.len() {
335                    features[i].clone()
336                } else {
337                    "*".to_string()
338                }
339            } else {
340                // Unknown word category: use unk feature string
341                let category_idx = i - unk_start_idx;
342                let char_def = &self.config.dict.character_definition;
343                let categories = char_def.categories();
344                if category_idx < categories.len() {
345                    let cat_name = &categories[category_idx];
346                    unk_categories
347                        .get(cat_name)
348                        .cloned()
349                        .unwrap_or_else(|| format!("UNK_{cat_name}"))
350                } else {
351                    format!("UNK_{}", category_idx)
352                }
353            }
354        };
355
356        let left_id_map = {
357            let mut id_to_feat: std::collections::HashMap<u32, String> =
358                std::collections::HashMap::new();
359            for (i, fs) in feature_sets.iter().enumerate() {
360                let lid = fs.left_id;
361                id_to_feat.entry(lid).or_insert_with(|| {
362                    let feat = get_feature_string(i);
363                    let (_, lfeature, _) = self.config.dictionary_rewriter.rewrite(&feat);
364                    lfeature
365                });
366            }
367            let mut entries: Vec<(u32, String)> = id_to_feat.into_iter().collect();
368            entries.sort_by_key(|&(id, _)| id);
369            entries
370        };
371        let right_id_map = {
372            let mut id_to_feat: std::collections::HashMap<u32, String> =
373                std::collections::HashMap::new();
374            for (i, fs) in feature_sets.iter().enumerate() {
375                let rid = fs.right_id;
376                id_to_feat.entry(rid).or_insert_with(|| {
377                    let feat = get_feature_string(i);
378                    let (_, _, rfeature) = self.config.dictionary_rewriter.rewrite(&feat);
379                    rfeature
380                });
381            }
382            let mut entries: Vec<(u32, String)> = id_to_feat.into_iter().collect();
383            entries.sort_by_key(|&(id, _)| id);
384            entries
385        };
386
387        let serializable_model = SerializableModel {
388            feature_weights,
389            labels: self.labels.clone(),
390            pos_info: self.extract_pos_info(),
391            metadata: ModelMetadata {
392                version: "1.0.0".to_string(),
393                regularization: self.regularization_cost,
394                iterations: self.max_iterations,
395                feature_count: self.feature_weights.len(),
396                label_count: self.labels.len(),
397            },
398            connection_matrix,
399            max_left_id,
400            max_right_id,
401            feature_sets,
402            unk_category_names,
403            unk_categories: self.config.unk_categories.clone(),
404            char_def_content: self.config.char_def_content.clone(),
405            feature_def_content: self.config.feature_def_content.clone(),
406            rewrite_def_content: self.config.rewrite_def_content.clone(),
407            cost_factor,
408            left_id_map,
409            right_id_map,
410        };
411
412        // Use rkyv for efficient binary serialization
413        let encoded = rkyv::to_bytes::<rkyv::rancor::Error>(&serializable_model)
414            .map_err(|e| anyhow::anyhow!("Failed to serialize model: {}", e))?;
415        writer.write_all(&encoded)?;
416
417        Ok(())
418    }
419
420    /// Reads a trained model from a reader.
421    ///
422    /// Accepts the current rkyv binary format and, for backward
423    /// compatibility, the legacy JSON format.
424    ///
425    /// # Arguments
426    ///
427    /// * `reader` - Reader containing the serialized model data
428    ///
429    /// # Returns
430    ///
431    /// The deserialized model.
432    ///
433    /// # Errors
434    ///
435    /// Returns an error naming the format that failed. A model written by a
436    /// build with a different `SerializableModel` layout is reported as such,
437    /// with the instruction to re-train -- it used to fall through to the
438    /// legacy JSON reader and surface as an unrelated UTF-8 error (#974).
439    pub fn read_model<R: Read>(mut reader: R) -> Result<SerializableModel> {
440        let mut buffer = Vec::new();
441        reader.read_to_end(&mut buffer)?;
442
443        if buffer.is_empty() {
444            return Err(anyhow::anyhow!(
445                "model file is empty; re-run `lindera train` to produce one"
446            ));
447        }
448
449        // Current format: rkyv. Tried first because it is self-validating, so
450        // a successful parse proves the on-disk layout matched this build.
451        let rkyv_error = match rkyv::from_bytes::<SerializableModel, rkyv::rancor::Error>(&buffer) {
452            Ok(mut model) => {
453                // Backward compatibility: if feature_sets is empty, generate from feature_weights
454                if model.feature_sets.is_empty() {
455                    model.feature_sets = model
456                        .feature_weights
457                        .iter()
458                        .map(|&weight| FeatureSetInfo {
459                            left_id: 0,
460                            right_id: 0,
461                            weight,
462                        })
463                        .collect();
464                }
465
466                return Ok(model);
467            }
468            Err(err) => err,
469        };
470
471        // Legacy format: JSON, attempted only when the payload actually looks
472        // like JSON. Routing on the leading byte as the *primary* check would
473        // be wrong -- a valid rkyv payload starts with the low byte of
474        // `feature_weights[0]` as a little-endian f64, which is `{` about one
475        // time in 256.
476        if looks_like_json_object(&buffer) {
477            let json_str = String::from_utf8(buffer).map_err(|err| {
478                anyhow::anyhow!("model file starts like legacy JSON but is not valid UTF-8: {err}")
479            })?;
480            return serde_json::from_str(&json_str)
481                .map_err(|err| anyhow::anyhow!("failed to parse legacy JSON model: {err}"));
482        }
483
484        Err(anyhow::anyhow!(
485            "failed to deserialize model: {rkyv_error}. The file is not a legacy JSON \
486             model either, so it was most likely written by a different version of \
487             lindera-trainer; re-run `lindera train` to regenerate it."
488        ))
489    }
490
491    /// Gets the merged model, creating it if necessary
492    fn get_merged_model(&self) -> Result<lindera_crf::MergedModel> {
493        Ok(self.raw_model.merge()?)
494    }
495
496    /// Generate a default feature string with appropriate number of fields
497    /// based on existing entries in the dictionary
498    fn generate_default_features(&self) -> String {
499        // Try to infer field count from existing unk_categories
500        if let Some(first_unk) = self.config.unk_categories.values().next() {
501            let field_count = first_unk.split(',').count();
502            return vec!["*"; field_count].join(",");
503        }
504
505        // Fallback: try from config.features
506        if let Some(first_feature) = self.config.features.first() {
507            let field_count = first_feature.split(',').count();
508            return vec!["*"; field_count].join(",");
509        }
510
511        // Ultimate fallback (should rarely happen)
512        "*".to_string()
513    }
514
515    /// Extracts part-of-speech information for each label
516    fn extract_pos_info(&self) -> Vec<String> {
517        // Get POS info from config.features (parallel to surfaces/labels)
518        let mut pos_info = Vec::new();
519
520        for (i, label) in self.labels.iter().enumerate() {
521            // First check if this is within the vocabulary (config.features range)
522            if i < self.config.features.len() {
523                pos_info.push(self.config.features[i].clone());
524            } else {
525                // For unknown word categories (DEFAULT, HIRAGANA, etc.),
526                // look up POS info from unk_categories
527                let unk_features = self
528                    .config
529                    .unk_categories
530                    .get(label)
531                    .cloned()
532                    .unwrap_or_else(|| {
533                        // Fallback: use DEFAULT category if label not found in unk_categories
534                        self.config
535                            .unk_categories
536                            .get("DEFAULT")
537                            .cloned()
538                            .unwrap_or_else(|| self.generate_default_features())
539                    });
540                pos_info.push(unk_features);
541            }
542        }
543
544        pos_info
545    }
546
547    /// Writes the dictionary files in Lindera format.
548    pub fn write_dictionary<W1, W2, W3, W4>(
549        &self,
550        lexicon_wtr: &mut W1,
551        connector_wtr: &mut W2,
552        unk_handler_wtr: &mut W3,
553        user_lexicon_wtr: &mut W4,
554    ) -> Result<()>
555    where
556        W1: Write,
557        W2: Write,
558        W3: Write,
559        W4: Write,
560    {
561        // Write lexicon with trained weights
562        self.write_lexicon(lexicon_wtr)?;
563
564        // Write connection cost matrix with trained costs
565        self.write_connection_costs(connector_wtr)?;
566
567        // Write unknown word handler with trained parameters
568        self.write_unknown_dictionary(unk_handler_wtr)?;
569
570        // Write user lexicon with trained weights
571        self.write_user_lexicon(user_lexicon_wtr)?;
572
573        Ok(())
574    }
575
576    pub fn write_lexicon<W: Write>(&self, writer: &mut W) -> Result<()> {
577        let merged_model = self.get_merged_model()?;
578        let cost_factor = calculate_cost_factor(&merged_model);
579
580        for (i, surface) in self.config.surfaces.iter().enumerate() {
581            if i < merged_model.feature_sets.len() {
582                let feature_set = merged_model.feature_sets[i];
583                let cost = tocost(feature_set.weight, cost_factor);
584                let features = &self.config.features[i];
585                writeln!(
586                    writer,
587                    "{},{},{},{},{}",
588                    surface,
589                    feature_set.left_id.get(),
590                    feature_set.right_id.get(),
591                    cost,
592                    features
593                )?;
594            } else {
595                let features = &self.config.features[i];
596                writeln!(writer, "{surface},0,0,0,{features}")?;
597            }
598        }
599
600        Ok(())
601    }
602
603    pub fn write_connection_costs<W: Write>(&self, writer: &mut W) -> Result<()> {
604        let merged_model = self.get_merged_model()?;
605        let cost_factor = calculate_cost_factor(&merged_model);
606
607        // Dense matrix dimensions: (right_size, left_size)
608        // +1 for BOS/EOS (id=0)
609        let right_size = merged_model.right_conn_to_left_feats.len() + 1;
610        let left_size = merged_model.left_conn_to_right_feats.len() + 1;
611        writeln!(writer, "{right_size} {left_size}")?;
612
613        // Dense matrix: output all (right_id, left_id) pairs
614        // Unseen pairs get maximum penalty cost to block unlearned transitions,
615        // forcing Viterbi to prefer paths through learned POS bigram connections.
616        for right_id in 0..right_size {
617            for left_id in 0..left_size {
618                let cost = if let Some(weight) = merged_model
619                    .matrix
620                    .get(right_id)
621                    .and_then(|hm| hm.get(&(left_id as u32)))
622                    .copied()
623                {
624                    tocost(weight, cost_factor)
625                } else {
626                    i16::MAX
627                };
628                writeln!(writer, "{right_id} {left_id} {cost}")?;
629            }
630        }
631
632        Ok(())
633    }
634
635    pub fn write_unknown_dictionary<W: Write>(&self, writer: &mut W) -> Result<()> {
636        let merged_model = self.get_merged_model()?;
637        let cost_factor = calculate_cost_factor(&merged_model);
638
639        let char_def = &self.config.dict.character_definition;
640        let category_names = char_def.categories();
641        let unk_start_idx = self.config.surfaces.len();
642
643        for (i, category_name) in category_names.iter().enumerate() {
644            let feature_set_idx = unk_start_idx + i;
645            if feature_set_idx < merged_model.feature_sets.len() {
646                let feature_set = merged_model.feature_sets[feature_set_idx];
647                let cost = tocost(feature_set.weight, cost_factor);
648
649                let default_features = self.generate_default_features();
650                let features = self
651                    .config
652                    .unk_categories
653                    .get(category_name.as_str())
654                    .map(|s| s.as_str())
655                    .unwrap_or(default_features.as_str());
656
657                writeln!(
658                    writer,
659                    "{},{},{},{},{}",
660                    category_name,
661                    feature_set.left_id.get(),
662                    feature_set.right_id.get(),
663                    cost,
664                    features
665                )?;
666            }
667        }
668
669        Ok(())
670    }
671
672    /// Writes the user lexicon with trained parameters.
673    ///
674    /// Entries are written in the order `read_user_lexicon` read them, so the
675    /// output is byte-reproducible and diffable against the input file (#981).
676    /// An entry whose left id, right id and cost are all zero receives
677    /// inferred context ids and a trained cost; any other entry is re-emitted
678    /// with its input parameters unchanged, per the contract documented on
679    /// `read_user_lexicon`.
680    ///
681    /// # 引数
682    ///
683    /// * `writer` - Write sink for the user lexicon CSV.
684    ///
685    /// # 戻り値
686    ///
687    /// `Ok(())` on success; writing nothing when no user lexicon was read.
688    fn write_user_lexicon<W: Write>(&self, writer: &mut W) -> Result<()> {
689        if self.user_entries.is_empty() {
690            return Ok(());
691        }
692
693        let merged_model = self.get_merged_model()?;
694        let cost_factor = calculate_cost_factor(&merged_model);
695
696        for (word, entry, _label_id) in &self.user_entries {
697            let surface = word.surface();
698            let features = word.feature();
699            if entry.left_id() == 0 && entry.right_id() == 0 && entry.word_cost() == 0 {
700                let (left_id, right_id) = self.infer_context_ids(surface, features);
701                let raw_cost = self.get_user_word_cost(surface) as f64 / 1000.0;
702                let cost = tocost(raw_cost, cost_factor);
703                writeln!(writer, "{surface},{left_id},{right_id},{cost},{features}")?;
704            } else {
705                writeln!(
706                    writer,
707                    "{},{},{},{},{}",
708                    surface,
709                    entry.left_id(),
710                    entry.right_id(),
711                    entry.word_cost(),
712                    features
713                )?;
714            }
715        }
716
717        Ok(())
718    }
719
720    fn get_user_word_cost(&self, _surface: &str) -> i32 {
721        // Return trained cost for user lexicon words
722        // Could be based on trained model weights
723        800 // Slightly lower cost than default for user words
724    }
725
726    /// Calculate unknown word cost based on trained feature weights using dynamic calculation
727    pub fn get_unknown_word_cost(&self, category: usize) -> i32 {
728        // Get category name from character definition
729        let char_def = &self.config.dict.character_definition;
730        let category_names = char_def.categories();
731
732        if category < category_names.len() {
733            let category_name = &category_names[category];
734            // Look up cost from unk_costs, with fallback to 2000
735            self.config
736                .unk_costs
737                .get(category_name)
738                .copied()
739                .unwrap_or(2000)
740        } else {
741            2000 // Default fallback cost
742        }
743    }
744
745    /// Infer context IDs (left_id, right_id) from surface form and features.
746    /// Finds the most similar POS pattern from trained vocabulary and uses its ID.
747    fn infer_context_ids(&self, surface: &str, features: &str) -> (u32, u32) {
748        // Parse feature string to get POS information
749        let feature_parts: Vec<&str> = features.split(',').collect();
750
751        // Find best matching entry from trained vocabulary by comparing features
752        // Try to match increasingly general patterns:
753        // 1. Exact feature match (all fields)
754        // 2. First 2 fields match (main POS + sub POS)
755        // 3. First field match (main POS only)
756        // 4. Same character category (via CharacterDefinition)
757
758        let char_def = &self.config.dict.character_definition;
759
760        // Strategy 1 & 2 & 3: Match by feature similarity
761        let mut best_match_idx: Option<usize> = None;
762        let mut best_match_score = 0;
763
764        for (i, _label) in self.labels.iter().enumerate() {
765            // Skip unknown word categories
766            if i >= self.config.features.len() {
767                break;
768            }
769
770            let vocab_features = &self.config.features[i];
771            let vocab_parts: Vec<&str> = vocab_features.split(',').collect();
772
773            // Calculate similarity score
774            let mut score = 0;
775            let max_fields = feature_parts.len().min(vocab_parts.len());
776
777            for j in 0..max_fields {
778                if feature_parts[j] == vocab_parts[j] {
779                    // Weight earlier fields more heavily (POS > sub-POS > details)
780                    score += (max_fields - j) * 10;
781                }
782            }
783
784            if score > best_match_score {
785                best_match_score = score;
786                best_match_idx = Some(i);
787            }
788        }
789
790        // If found a match from vocabulary, look up from user_entries
791        if let Some(idx) = best_match_idx {
792            // Try to find in user_entries first
793            if idx < self.user_entries.len() {
794                let (_, entry, _) = &self.user_entries[idx];
795                return (entry.left_id(), entry.right_id());
796            }
797        }
798
799        // Strategy 4: If no good match found, use character category
800        if best_match_score == 0 && !surface.is_empty() {
801            let first_char = surface.chars().next().unwrap();
802            let categories = char_def.lookup_categories(first_char);
803
804            if !categories.is_empty() {
805                let category_id = categories[0].0 as u32;
806                // Use category ID as both left and right ID
807                return (category_id, category_id);
808            }
809        }
810
811        // Ultimate fallback: use first user_entry's IDs or default to 0
812        if let Some((_, entry, _)) = self.user_entries.first() {
813            (entry.left_id(), entry.right_id())
814        } else {
815            (0, 0)
816        }
817    }
818
819    /// Calculate maximum context ID from training data
820    #[allow(dead_code)]
821    fn calculate_max_context_id(&self) -> u32 {
822        let mut max_id = 0u32;
823
824        // Get maximum ID from all user_entries
825        for (_, entry, _) in &self.user_entries {
826            max_id = max_id.max(entry.left_id()).max(entry.right_id());
827        }
828
829        max_id
830    }
831
832    /// Calculate connection cost based on trained model
833    #[allow(dead_code)]
834    fn get_trained_connection_cost(&self, from_id: usize, to_id: usize) -> i32 {
835        // Use CRF feature weights to calculate connection cost
836        let weights = self.raw_model.weights();
837
838        if weights.is_empty() {
839            return 0; // Fallback
840        }
841
842        // Calculate cost based on context ID combination
843        let cost_index = (from_id * 1000 + to_id) % weights.len();
844        let raw_cost = weights[cost_index];
845
846        // Negative weight = low cost (good connection), positive weight = high cost (bad connection)
847        let scaled_cost = (-raw_cost * 1000.0) as i32;
848
849        // Limit cost range (-10000 to 10000)
850        scaled_cost.clamp(-10000, 10000)
851    }
852
853    /// Gets the number of features in the model.
854    pub fn num_features(&self) -> usize {
855        // Return the actual feature count from the raw model
856        // The raw model contains the feature weights vector
857        self.feature_weights.len()
858    }
859
860    /// Gets the number of labels in the model.
861    pub fn num_labels(&self) -> usize {
862        // Return the actual label count
863        self.labels.len()
864    }
865
866    /// Gets the raw CRF model for advanced operations.
867    pub fn raw_model(&self) -> &lindera_crf::RawModel {
868        &self.raw_model
869    }
870
871    /// Writes the bigram details in three separate files.
872    ///
873    /// This method outputs:
874    /// - Left features: per left connection id, the quoted right-side (`%R`)
875    ///   feature strings it groups, `*` for absent slots
876    /// - Right features: per right connection id, the quoted left-side (`%L`)
877    ///   feature strings it groups
878    /// - Costs: one `{left}/{right}\t{cost}` line per raw bigram feature
879    ///   pair, where id 0 renders as `BOS` / `EOS` and names come from the
880    ///   feature extractor
881    ///
882    /// All three files are byte-reproducible for the same model: the loops
883    /// run over vectors or key-sorted map entries, never over raw map
884    /// iteration order (#981).
885    ///
886    /// # 引数
887    ///
888    /// * `left_wtr` - Write sink for the left-feature file.
889    /// * `right_wtr` - Write sink for the right-feature file.
890    /// * `cost_wtr` - Write sink for the feature-pair cost file.
891    ///
892    /// # 戻り値
893    ///
894    /// `Ok(())` on success.
895    pub fn write_bigram_details<L, R, C>(
896        &self,
897        left_wtr: L,
898        right_wtr: R,
899        cost_wtr: C,
900    ) -> Result<()>
901    where
902        L: Write,
903        R: Write,
904        C: Write,
905    {
906        use std::collections::HashMap;
907        use std::io::BufWriter;
908
909        // Get merged model for detailed analysis
910        let merged_model = self.get_merged_model()?;
911        let cost_factor = self.config.cost_factor;
912
913        // Real feature names, inverted from the extractor's string-to-id maps
914        // (the same inversion write_left_id_def uses): id -> feature string.
915        let right_features: HashMap<u32, &str> = self
916            .config
917            .feature_extractor
918            .right_feature_ids
919            .iter()
920            .map(|(s, &id)| (id.get(), s.as_str()))
921            .collect();
922        let left_features: HashMap<u32, &str> = self
923            .config
924            .feature_extractor
925            .left_feature_ids
926            .iter()
927            .map(|(s, &id)| (id.get(), s.as_str()))
928            .collect();
929
930        // Write left features
931        let mut left_wtr = BufWriter::new(left_wtr);
932        for (conn_id, left_feat_ids) in merged_model.left_conn_to_right_feats.iter().enumerate() {
933            write!(&mut left_wtr, "{}\t", conn_id + 1)?;
934            for (i, feat_id) in left_feat_ids.iter().enumerate() {
935                if i != 0 {
936                    write!(&mut left_wtr, ",")?;
937                }
938                if let Some(feat_id) = feat_id {
939                    let feat_str = right_features.get(&feat_id.get()).copied().unwrap_or("*");
940                    write!(&mut left_wtr, "\"{feat_str}\"")?;
941                } else {
942                    write!(&mut left_wtr, "*")?;
943                }
944            }
945            writeln!(&mut left_wtr)?;
946        }
947
948        // Write right features
949        let mut right_wtr = BufWriter::new(right_wtr);
950        for (conn_id, right_feat_ids) in merged_model.right_conn_to_left_feats.iter().enumerate() {
951            write!(&mut right_wtr, "{}\t", conn_id + 1)?;
952            for (i, feat_id) in right_feat_ids.iter().enumerate() {
953                if i != 0 {
954                    write!(&mut right_wtr, ",")?;
955                }
956                if let Some(feat_id) = feat_id {
957                    let feat_str = left_features.get(&feat_id.get()).copied().unwrap_or("*");
958                    write!(&mut right_wtr, "\"{feat_str}\"")?;
959                } else {
960                    write!(&mut right_wtr, "*")?;
961                }
962            }
963            writeln!(&mut right_wtr)?;
964        }
965
966        // Write bigram costs with feature pair names. The source is the raw
967        // model's bigram weight table -- left bigram feature id (0 = BOS) to
968        // right bigram feature id (0 = EOS) to weight index -- which is the
969        // id space these name maps describe. The inner maps are key-sorted so
970        // the output is byte-reproducible (#981).
971        let mut cost_wtr = BufWriter::new(cost_wtr);
972        for (left_feat_id, hm) in self.raw_model.bigram_weight_indices().iter().enumerate() {
973            let left_feat_str = if left_feat_id == 0 {
974                "BOS"
975            } else {
976                left_features
977                    .get(&(left_feat_id as u32))
978                    .copied()
979                    .unwrap_or("*")
980            };
981            let mut entries: Vec<(u32, u32)> = hm.iter().map(|(&k, &v)| (k, v)).collect();
982            entries.sort_unstable_by_key(|&(k, _)| k);
983            for (right_feat_id, weight_index) in entries {
984                let right_feat_str = if right_feat_id == 0 {
985                    "EOS"
986                } else {
987                    right_features.get(&right_feat_id).copied().unwrap_or("*")
988                };
989                let weight = self
990                    .raw_model
991                    .weights()
992                    .get(weight_index as usize)
993                    .copied()
994                    .ok_or_else(|| {
995                        anyhow::anyhow!("bigram weight index {weight_index} out of range")
996                    })?;
997                let cost = tocost(weight, cost_factor);
998                writeln!(&mut cost_wtr, "{left_feat_str}/{right_feat_str}\t{cost}")?;
999            }
1000        }
1001
1002        Ok(())
1003    }
1004
1005    /// Evaluates the model on test data.
1006    /// Returns a simple evaluation score based on feature weights.
1007    pub fn evaluate(&self, _test_lattices: &[lindera_crf::Lattice]) -> f64 {
1008        // For now, return a simple evaluation based on the model's feature weights
1009        // A more sophisticated implementation would compute actual likelihood scores
1010        let weights = self.raw_model.weights();
1011
1012        // Compute average absolute weight as a simple evaluation metric
1013        if weights.is_empty() {
1014            0.0
1015        } else {
1016            let sum: f64 = weights.iter().map(|w| w.abs()).sum();
1017            sum / weights.len() as f64
1018        }
1019    }
1020
1021    /// Write dictionary components to separate buffers using optimized serialization
1022    pub fn write_dictionary_buffers(
1023        &self,
1024        lexicon: &mut Vec<u8>,
1025        connector: &mut Vec<u8>,
1026        unk_handler: &mut Vec<u8>,
1027        user_lexicon: &mut Vec<u8>,
1028    ) -> Result<()> {
1029        // Serialize lexicon data
1030        let lexicon_data = rkyv::to_bytes::<rkyv::rancor::Error>(&self.labels)
1031            .map_err(|e| anyhow::anyhow!("Failed to serialize lexicon: {}", e))?;
1032        lexicon.extend_from_slice(&lexicon_data);
1033
1034        // Serialize connection costs (feature weights as connection matrix)
1035        let connection_data = rkyv::to_bytes::<rkyv::rancor::Error>(&self.feature_weights)
1036            .map_err(|e| anyhow::anyhow!("Failed to serialize connector: {}", e))?;
1037        connector.extend_from_slice(&connection_data);
1038
1039        // Serialize unknown word handler (simplified data)
1040        let unk_data = rkyv::to_bytes::<rkyv::rancor::Error>(&self.user_entries.len())
1041            .map_err(|e| anyhow::anyhow!("Failed to serialize unknown handler: {}", e))?;
1042        unk_handler.extend_from_slice(&unk_data);
1043
1044        // Serialize user lexicon (config info as user lexicon)
1045        let user_data = rkyv::to_bytes::<rkyv::rancor::Error>(&self.config.surfaces)
1046            .map_err(|e| anyhow::anyhow!("Failed to serialize user lexicon: {}", e))?;
1047        user_lexicon.extend_from_slice(&user_data);
1048
1049        Ok(())
1050    }
1051}
1052
1053impl SerializableModel {
1054    /// Generate a default feature string with appropriate number of fields
1055    /// based on existing entries in the dictionary
1056    fn generate_default_features(&self) -> String {
1057        // Try to infer field count from existing unk_categories
1058        if let Some(first_unk) = self.unk_categories.values().next() {
1059            let field_count = first_unk.split(',').count();
1060            return vec!["*"; field_count].join(",");
1061        }
1062
1063        // Fallback: try from pos_info
1064        if let Some(first_pos) = self.pos_info.first() {
1065            let field_count = first_pos.split(',').count();
1066            return vec!["*"; field_count].join(",");
1067        }
1068
1069        // Ultimate fallback (should rarely happen)
1070        "*".to_string()
1071    }
1072
1073    /// Write lexicon file with proper cost calculation
1074    pub fn write_lexicon<W: std::io::Write>(&self, writer: &mut W) -> anyhow::Result<()> {
1075        // Unknown word categories are at the end of labels, skip them (they go to unk.def)
1076        let unk_start_idx = self
1077            .labels
1078            .len()
1079            .saturating_sub(self.unk_category_names.len());
1080
1081        // Write lexicon entries using learned connection IDs and costs
1082        for (i, label) in self.labels.iter().enumerate() {
1083            if i >= unk_start_idx {
1084                continue;
1085            }
1086
1087            if i < self.feature_sets.len() {
1088                let fs = &self.feature_sets[i];
1089                let pos_info_str;
1090                let pos_info = if i < self.pos_info.len() {
1091                    &self.pos_info[i]
1092                } else {
1093                    pos_info_str = self.generate_default_features();
1094                    &pos_info_str
1095                };
1096
1097                let cost = tocost(fs.weight, self.cost_factor);
1098
1099                writeln!(
1100                    writer,
1101                    "{},{},{},{},{}",
1102                    label, fs.left_id, fs.right_id, cost, pos_info
1103                )?;
1104            }
1105        }
1106
1107        Ok(())
1108    }
1109
1110    /// Write dense connection cost matrix (MeCab-compatible)
1111    pub fn write_connection_costs<W: std::io::Write>(&self, writer: &mut W) -> anyhow::Result<()> {
1112        if !self.connection_matrix.is_empty() {
1113            let right_size = self.max_right_id + 1;
1114            let left_size = self.max_left_id + 1;
1115
1116            writeln!(writer, "{right_size} {left_size}")?;
1117
1118            // Dense matrix: all (right_id, left_id) pairs
1119            // Unseen pairs get maximum penalty cost to block unlearned transitions,
1120            // forcing Viterbi to prefer paths through learned POS bigram connections.
1121            for right_id in 0..right_size {
1122                for left_id in 0..left_size {
1123                    let cost = if let Some(&weight) = self
1124                        .connection_matrix
1125                        .get(&right_id)
1126                        .and_then(|inner| inner.get(&left_id))
1127                    {
1128                        tocost(weight, self.cost_factor)
1129                    } else {
1130                        i16::MAX
1131                    };
1132                    writeln!(writer, "{right_id} {left_id} {cost}")?;
1133                }
1134            }
1135        } else {
1136            writeln!(writer, "0 0")?;
1137        }
1138
1139        Ok(())
1140    }
1141
1142    /// Update metadata.json with trained model values.
1143    ///
1144    /// The emitted `model_info` section deliberately carries no timestamp, so
1145    /// exporting the same model twice produces byte-identical output (#981);
1146    /// serde_json's map is ordered, so key order is stable too.
1147    ///
1148    /// # 引数
1149    ///
1150    /// * `base_metadata_path` - Path of the base `metadata.json` to update.
1151    /// * `writer` - Write sink for the updated JSON.
1152    pub fn update_metadata_json<W: std::io::Write>(
1153        &self,
1154        base_metadata_path: &std::path::Path,
1155        writer: &mut W,
1156    ) -> anyhow::Result<()> {
1157        // Read the base metadata.json file
1158        let base_content = std::fs::read_to_string(base_metadata_path)?;
1159        let mut metadata: serde_json::Value = serde_json::from_str(&base_content)?;
1160
1161        // Calculate updated values based on trained model
1162        let updated_default_cost = if !self.feature_weights.is_empty() {
1163            // Calculate median feature weight for default cost
1164            let mut weights = self.feature_weights.clone();
1165            weights.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
1166            let median_weight = if weights.is_empty() {
1167                0.0
1168            } else {
1169                weights[weights.len() / 2]
1170            };
1171            // Convert to appropriate cost range for practical use
1172            (median_weight * 500.0).abs() as i32 + 1500
1173        } else {
1174            // Keep existing value if no trained weights
1175            metadata
1176                .get("default_word_cost")
1177                .and_then(|v| v.as_i64())
1178                .unwrap_or(2000) as i32
1179        };
1180
1181        // Update metadata with trained model values
1182        metadata["default_word_cost"] =
1183            serde_json::Value::Number(serde_json::Number::from(updated_default_cost));
1184
1185        // Add model_info section with training statistics
1186        let max_context_id = std::cmp::max(self.max_left_id, self.max_right_id);
1187        metadata["model_info"] = serde_json::json!({
1188            "feature_count": self.feature_weights.len(),
1189            "label_count": self.labels.len(),
1190            "max_left_context_id": self.max_left_id,
1191            "max_right_context_id": self.max_right_id,
1192            "connection_matrix_size": format!("{}x{}", max_context_id + 1, max_context_id + 1),
1193            "version": self.metadata.version,
1194            "training_iterations": self.metadata.iterations,
1195            "regularization": self.metadata.regularization,
1196        });
1197
1198        // Write updated metadata
1199        let formatted = serde_json::to_string_pretty(&metadata)?;
1200        writer.write_all(formatted.as_bytes())?;
1201
1202        Ok(())
1203    }
1204
1205    /// Write unknown word definitions
1206    pub fn write_unknown_dictionary<W: std::io::Write>(
1207        &self,
1208        writer: &mut W,
1209    ) -> anyhow::Result<()> {
1210        let unk_start_idx = self
1211            .labels
1212            .len()
1213            .saturating_sub(self.unk_category_names.len());
1214
1215        for (i, category_name) in self.unk_category_names.iter().enumerate() {
1216            let feature_idx = unk_start_idx + i;
1217
1218            if feature_idx < self.feature_sets.len() {
1219                let fs = &self.feature_sets[feature_idx];
1220
1221                let features = self
1222                    .unk_categories
1223                    .get(category_name)
1224                    .cloned()
1225                    .unwrap_or_else(|| self.generate_default_features());
1226
1227                let cost = tocost(fs.weight, self.cost_factor);
1228
1229                writeln!(
1230                    writer,
1231                    "{},{},{},{},{}",
1232                    category_name, fs.left_id, fs.right_id, cost, features
1233                )?;
1234            }
1235        }
1236
1237        Ok(())
1238    }
1239
1240    /// Writes the character definition file (char.def) content preserved from training.
1241    pub fn write_char_def<W: std::io::Write>(&self, writer: &mut W) -> anyhow::Result<()> {
1242        writer.write_all(self.char_def_content.as_bytes())?;
1243        Ok(())
1244    }
1245
1246    /// Writes the feature definition file (feature.def) content preserved from training.
1247    pub fn write_feature_def<W: std::io::Write>(&self, writer: &mut W) -> anyhow::Result<()> {
1248        writer.write_all(self.feature_def_content.as_bytes())?;
1249        Ok(())
1250    }
1251
1252    /// Writes the rewrite rule definition file (rewrite.def) content preserved from training.
1253    pub fn write_rewrite_def<W: std::io::Write>(&self, writer: &mut W) -> anyhow::Result<()> {
1254        writer.write_all(self.rewrite_def_content.as_bytes())?;
1255        Ok(())
1256    }
1257
1258    /// Writes left-id.def: maps left context IDs to their feature strings.
1259    pub fn write_left_id_def<W: std::io::Write>(&self, writer: &mut W) -> anyhow::Result<()> {
1260        writeln!(writer, "0 BOS/EOS")?;
1261        for (id, feat_str) in &self.left_id_map {
1262            writeln!(writer, "{id} {feat_str}")?;
1263        }
1264        Ok(())
1265    }
1266
1267    /// Writes right-id.def: maps right context IDs to their feature strings.
1268    pub fn write_right_id_def<W: std::io::Write>(&self, writer: &mut W) -> anyhow::Result<()> {
1269        writeln!(writer, "0 BOS/EOS")?;
1270        for (id, feat_str) in &self.right_id_map {
1271            writeln!(writer, "{id} {feat_str}")?;
1272        }
1273        Ok(())
1274    }
1275}
1276
1277impl Model {
1278    /// Writes left-id.def: maps left context IDs to their feature strings.
1279    ///
1280    /// Format: `id feature_string` (one per line), with BOS/EOS at id=0.
1281    pub fn write_left_id_def<W: Write>(&self, writer: &mut W) -> Result<()> {
1282        // id=0 is reserved for BOS/EOS
1283        writeln!(writer, "0 BOS/EOS")?;
1284
1285        // Build reverse mapping: NonZeroU32 -> feature string
1286        let left_ids = &self.config.feature_extractor.left_feature_ids;
1287        let mut entries: Vec<(u32, &str)> = left_ids
1288            .iter()
1289            .map(|(feat_str, &id)| (id.get(), feat_str.as_str()))
1290            .collect();
1291        entries.sort_by_key(|&(id, _)| id);
1292
1293        for (id, feat_str) in entries {
1294            writeln!(writer, "{id} {feat_str}")?;
1295        }
1296
1297        Ok(())
1298    }
1299
1300    /// Writes right-id.def: maps right context IDs to their feature strings.
1301    ///
1302    /// Format: `id feature_string` (one per line), with BOS/EOS at id=0.
1303    pub fn write_right_id_def<W: Write>(&self, writer: &mut W) -> Result<()> {
1304        // id=0 is reserved for BOS/EOS
1305        writeln!(writer, "0 BOS/EOS")?;
1306
1307        // Build reverse mapping: NonZeroU32 -> feature string
1308        let right_ids = &self.config.feature_extractor.right_feature_ids;
1309        let mut entries: Vec<(u32, &str)> = right_ids
1310            .iter()
1311            .map(|(feat_str, &id)| (id.get(), feat_str.as_str()))
1312            .collect();
1313        entries.sort_by_key(|&(id, _)| id);
1314
1315        for (id, feat_str) in entries {
1316            writeln!(writer, "{id} {feat_str}")?;
1317        }
1318
1319        Ok(())
1320    }
1321}
1322
1323#[cfg(test)]
1324mod tests {
1325    use crate::{Trainer, TrainerConfig};
1326    use std::io::Cursor;
1327
1328    #[test]
1329    fn test_trainer_creation() {
1330        // Test that Trainer can be created from a valid config
1331        let lexicon_data = "外国,0,0,5000,名詞,一般,*,*,*,*,外国,ガイコク,ガイコク\n";
1332        let char_data = "# char.def placeholder\n";
1333        let unk_data = "# unk.def placeholder\n";
1334        let feature_data = "UNIGRAM:%F[0]\nLEFT:%L[0]\nRIGHT:%R[0]\n";
1335        let rewrite_data = "# rewrite.def placeholder\n";
1336
1337        let config_result = TrainerConfig::from_readers(
1338            Cursor::new(lexicon_data.as_bytes()),
1339            Cursor::new(char_data.as_bytes()),
1340            Cursor::new(unk_data.as_bytes()),
1341            Cursor::new(feature_data.as_bytes()),
1342            Cursor::new(rewrite_data.as_bytes()),
1343        );
1344
1345        assert!(config_result.is_ok());
1346        let config = config_result.unwrap();
1347
1348        // Test trainer creation with builder pattern
1349        let trainer = Trainer::new(config)
1350            .unwrap()
1351            .regularization_cost(0.01)
1352            .max_iter(10)
1353            .num_threads(1);
1354
1355        // Verify trainer settings using the getters
1356        assert_eq!(trainer.get_regularization_cost(), 0.01);
1357        assert_eq!(trainer.get_max_iter(), 10);
1358        assert_eq!(trainer.get_num_threads(), 1);
1359    }
1360
1361    #[test]
1362    fn test_tocost() {
1363        use super::tocost;
1364
1365        // MeCab: tocost(d, n) = clamp(-n * d, -32767, 32767)
1366        // Positive weight → negative cost
1367        assert_eq!(tocost(1.0, 700), -700);
1368        // Negative weight → positive cost
1369        assert_eq!(tocost(-1.0, 700), 700);
1370        // Zero weight → zero cost
1371        assert_eq!(tocost(0.0, 700), 0);
1372        // Clamp to i16::MAX
1373        assert_eq!(tocost(-100.0, 700), i16::MAX);
1374        // Clamp to i16::MIN
1375        assert_eq!(tocost(100.0, 700), i16::MIN);
1376        // Fractional weight
1377        assert_eq!(tocost(0.5, 700), -350);
1378    }
1379}
1380
1381// Regression tests for #974: `model.dat` must be byte-reproducible.
1382//
1383// Feature ids and connection costs are stable across runs, but the archived
1384// bytes were not: rkyv lays an archived `HashMap` out as a SwissTable whose
1385// slot assignment resolves collisions in source iteration order, and
1386// `HashMap` iteration order is seeded per instance.
1387#[cfg(test)]
1388mod determinism_tests {
1389    use super::{FeatureSetInfo, ModelMetadata, SerializableModel};
1390    use std::collections::BTreeMap;
1391
1392    /// Builds a model whose two order-sensitive maps are populated by walking
1393    /// `keys` in the given order. Everything else is fixed.
1394    fn model_with_maps(keys: &[usize]) -> SerializableModel {
1395        let mut connection_matrix = BTreeMap::new();
1396        let mut unk_categories = BTreeMap::new();
1397        for &k in keys {
1398            // The inner key set is derived from `k`, not from the traversal
1399            // order, so two builds over the same keys in different orders
1400            // produce logically identical maps -- which is exactly what this
1401            // test needs to isolate ordering from content.
1402            let mut inner = BTreeMap::new();
1403            for j in 0..4 {
1404                inner.insert(k + j, (k * 64 + j) as f64);
1405            }
1406            connection_matrix.insert(k, inner);
1407            unk_categories.insert(format!("CAT{k:03}"), format!("feat{k},*,*"));
1408        }
1409        SerializableModel {
1410            feature_weights: vec![0.5; 4],
1411            labels: vec!["a".to_string()],
1412            pos_info: vec!["名詞,一般".to_string()],
1413            metadata: ModelMetadata {
1414                version: "1.0.0".to_string(),
1415                regularization: 0.01,
1416                iterations: 10,
1417                feature_count: 4,
1418                label_count: 1,
1419            },
1420            connection_matrix,
1421            max_left_id: 63,
1422            max_right_id: 63,
1423            feature_sets: vec![FeatureSetInfo {
1424                left_id: 1,
1425                right_id: 1,
1426                weight: 0.25,
1427            }],
1428            unk_category_names: vec!["CAT000".to_string()],
1429            unk_categories,
1430            char_def_content: String::new(),
1431            feature_def_content: String::new(),
1432            rewrite_def_content: String::new(),
1433            cost_factor: 700,
1434            left_id_map: vec![(1, "*".to_string())],
1435            right_id_map: vec![(1, "*".to_string())],
1436        }
1437    }
1438
1439    /// The archived bytes must depend only on the logical contents of the two
1440    /// maps, never on the order they were populated in.
1441    ///
1442    /// 32 keys is deliberate: rkyv sizes the SwissTable at ~89% load and its
1443    /// probe window is 16 slots, so under `HashMap` the ascending and
1444    /// descending builds necessarily resolve collisions differently. Two
1445    /// separately constructed `HashMap`s also get different `RandomState`
1446    /// seeds, so this is not a coin flip.
1447    #[test]
1448    fn serialized_model_is_independent_of_map_insertion_order() {
1449        let forward: Vec<usize> = (0..32).collect();
1450        let reverse: Vec<usize> = (0..32).rev().collect();
1451
1452        let a = match rkyv::to_bytes::<rkyv::rancor::Error>(&model_with_maps(&forward)) {
1453            Ok(bytes) => bytes,
1454            Err(err) => panic!("failed to serialize ascending model: {err}"),
1455        };
1456        let b = match rkyv::to_bytes::<rkyv::rancor::Error>(&model_with_maps(&reverse)) {
1457            Ok(bytes) => bytes,
1458            Err(err) => panic!("failed to serialize descending model: {err}"),
1459        };
1460
1461        assert_eq!(
1462            a.as_slice(),
1463            b.as_slice(),
1464            "archived bytes must not depend on map insertion order"
1465        );
1466    }
1467
1468    /// Serializing the very same value twice must also be stable -- a weaker
1469    /// property than the one above, but it is the one a caller writing the
1470    /// same model to two files relies on.
1471    #[test]
1472    fn serializing_the_same_model_twice_is_stable() {
1473        let keys: Vec<usize> = (0..32).collect();
1474        let model = model_with_maps(&keys);
1475
1476        let a = match rkyv::to_bytes::<rkyv::rancor::Error>(&model) {
1477            Ok(bytes) => bytes,
1478            Err(err) => panic!("first serialization failed: {err}"),
1479        };
1480        let b = match rkyv::to_bytes::<rkyv::rancor::Error>(&model) {
1481            Ok(bytes) => bytes,
1482            Err(err) => panic!("second serialization failed: {err}"),
1483        };
1484
1485        assert_eq!(a.as_slice(), b.as_slice());
1486    }
1487
1488    /// The connection matrix's type, spelled out so the compile-time guard
1489    /// below names it exactly.
1490    type OrderedConnectionMatrix = BTreeMap<usize, BTreeMap<usize, f64>>;
1491
1492    /// The unknown-word category map's type, for the same reason.
1493    type OrderedUnkCategories = BTreeMap<String, String>;
1494
1495    /// Compile-time guard. `model.dat` is byte-reproducible only while these
1496    /// two fields are ordered maps. These functions are never called; they
1497    /// exist so that switching either field back to `HashMap` fails to
1498    /// compile rather than silently regressing (#974).
1499    #[allow(dead_code)]
1500    fn assert_connection_matrix_is_ordered(model: &SerializableModel) -> &OrderedConnectionMatrix {
1501        &model.connection_matrix
1502    }
1503
1504    /// See [`assert_connection_matrix_is_ordered`].
1505    #[allow(dead_code)]
1506    fn assert_unk_categories_is_ordered(model: &SerializableModel) -> &OrderedUnkCategories {
1507        &model.unk_categories
1508    }
1509
1510    /// A payload that is neither valid rkyv nor JSON must say so and tell the
1511    /// user to re-train, rather than surfacing an unrelated UTF-8 error.
1512    #[test]
1513    fn read_model_rejects_unrecognized_payload_with_an_actionable_error() {
1514        let garbage = vec![0xffu8; 64];
1515
1516        let err = match super::Model::read_model(garbage.as_slice()) {
1517            Ok(_) => panic!("garbage must not deserialize"),
1518            Err(err) => err.to_string(),
1519        };
1520
1521        assert!(err.contains("lindera train"), "unhelpful message: {err}");
1522    }
1523
1524    /// A payload that looks like JSON but is malformed is reported as a JSON
1525    /// parse failure, so the legacy path stays diagnosable.
1526    #[test]
1527    fn read_model_reports_malformed_legacy_json() {
1528        let err = match super::Model::read_model(&b"{ not json"[..]) {
1529            Ok(_) => panic!("malformed JSON must not deserialize"),
1530            Err(err) => err.to_string(),
1531        };
1532
1533        assert!(err.contains("legacy JSON"), "unhelpful message: {err}");
1534    }
1535
1536    /// An empty file is its own case, and used to reach the JSON reader.
1537    #[test]
1538    fn read_model_rejects_an_empty_file() {
1539        let err = match super::Model::read_model(&b""[..]) {
1540            Ok(_) => panic!("an empty file must not deserialize"),
1541            Err(err) => err.to_string(),
1542        };
1543
1544        assert!(err.contains("empty"), "unhelpful message: {err}");
1545    }
1546
1547    /// `update_metadata_json` must be byte-reproducible and must not emit a
1548    /// timestamp; its output must round-trip through
1549    /// `lindera_dictionary`'s `Metadata` loader.
1550    ///
1551    /// Before #981 it stamped `updated_at` from `SystemTime::now()` (and
1552    /// `unwrap()`ed in production code), so two exports of the same model
1553    /// could never compare equal.
1554    #[test]
1555    fn metadata_json_is_reproducible_and_timestamp_free() {
1556        use lindera_dictionary::dictionary::metadata::Metadata;
1557
1558        let model = model_with_maps(&[0, 7, 21]);
1559
1560        let dir = match tempfile::tempdir() {
1561            Ok(dir) => dir,
1562            Err(err) => panic!("failed to create temp dir: {err}"),
1563        };
1564        let base_path = dir.path().join("metadata.json");
1565        let base_json = match serde_json::to_string_pretty(&Metadata::default()) {
1566            Ok(json) => json,
1567            Err(err) => panic!("failed to serialize base metadata: {err}"),
1568        };
1569        if let Err(err) = std::fs::write(&base_path, base_json) {
1570            panic!("failed to write base metadata: {err}");
1571        }
1572
1573        let update = || -> Vec<u8> {
1574            let mut buf = Vec::new();
1575            if let Err(err) = model.update_metadata_json(&base_path, &mut buf) {
1576                panic!("update_metadata_json failed: {err}");
1577            }
1578            buf
1579        };
1580
1581        let first = update();
1582        let second = update();
1583        assert_eq!(first, second, "metadata.json output must be reproducible");
1584
1585        let value: serde_json::Value = match serde_json::from_slice(&first) {
1586            Ok(value) => value,
1587            Err(err) => panic!("output is not valid JSON: {err}"),
1588        };
1589        let model_info = match value.get("model_info") {
1590            Some(serde_json::Value::Object(map)) => map,
1591            other => panic!("model_info must be an object, got {other:?}"),
1592        };
1593        assert!(
1594            !model_info.contains_key("updated_at"),
1595            "model_info must not carry a timestamp"
1596        );
1597        assert!(model_info.contains_key("feature_count"));
1598        assert!(model_info.contains_key("version"));
1599
1600        // The emitted file must load through the dictionary-side struct,
1601        // pinning writer/struct consistency end to end.
1602        if let Err(err) = Metadata::load(&first) {
1603            panic!("exported metadata.json must load as Metadata: {err}");
1604        }
1605    }
1606}