lindera_trainer/
corpus.rs

1use std::io::{BufRead, BufReader, Read};
2
3use anyhow::Result;
4
5/// Representation of a pair of a surface and features.
6pub struct Word {
7    surface: String,
8    feature: String,
9}
10
11impl Word {
12    pub fn new(surface: &str, feature: &str) -> Self {
13        Self {
14            surface: surface.to_string(),
15            feature: feature.to_string(),
16        }
17    }
18
19    /// Returns a surface string.
20    pub fn surface(&self) -> &str {
21        &self.surface
22    }
23
24    /// Returns a concatenated feature string.
25    pub fn feature(&self) -> &str {
26        &self.feature
27    }
28}
29
30/// Representation of a sentence.
31pub struct Example {
32    pub(crate) sentence: String,
33    pub(crate) tokens: Vec<Word>,
34}
35
36impl Example {
37    /// Creates a new example from tokens.
38    pub fn new(tokens: Vec<Word>) -> Self {
39        let sentence = tokens.iter().map(|t| t.surface()).collect::<String>();
40        Self { sentence, tokens }
41    }
42}
43
44/// A corpus for training.
45pub struct Corpus {
46    pub(crate) examples: Vec<Example>,
47}
48
49impl Default for Corpus {
50    fn default() -> Self {
51        Self::new()
52    }
53}
54
55impl Corpus {
56    /// Creates a new empty corpus.
57    pub fn new() -> Self {
58        Self {
59            examples: Vec::new(),
60        }
61    }
62
63    /// Creates a corpus from a reader.
64    ///
65    /// The format should be compatible with standard tokenized output:
66    /// Each line contains a token with tab-separated surface and features.
67    /// Sentences are separated by empty lines.
68    pub fn from_reader<R: Read>(reader: R) -> Result<Self> {
69        let reader = BufReader::new(reader);
70        let mut corpus = Self::new();
71        let mut current_tokens = Vec::new();
72
73        for line in reader.lines() {
74            let line = line?;
75
76            if line.trim().is_empty() || line == "EOS" {
77                // End of sentence
78                if !current_tokens.is_empty() {
79                    corpus.examples.push(Example::new(current_tokens));
80                    current_tokens = Vec::new();
81                }
82            } else {
83                // Parse token line: surface\tfeatures
84                let parts: Vec<&str> = line.split('\t').collect();
85                if parts.len() >= 2 {
86                    let surface = parts[0];
87                    let features = parts[1];
88                    current_tokens.push(Word::new(surface, features));
89                }
90            }
91        }
92
93        // Add remaining tokens if any
94        if !current_tokens.is_empty() {
95            corpus.examples.push(Example::new(current_tokens));
96        }
97
98        Ok(corpus)
99    }
100
101    /// Returns the number of examples in the corpus.
102    pub fn len(&self) -> usize {
103        self.examples.len()
104    }
105
106    /// Returns true if the corpus is empty.
107    pub fn is_empty(&self) -> bool {
108        self.examples.is_empty()
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use std::io::Cursor;
116
117    #[test]
118    fn test_corpus_from_reader() {
119        let corpus_data = r#"外国	名詞,一般,*,*,*,*,外国,ガイコク,ガイコク
120人	名詞,接尾,一般,*,*,*,人,ジン,ジン
121EOS
122
123これ	連体詞,*,*,*,*,*,これ,コレ,コレ
124は	助詞,係助詞,*,*,*,*,は,ハ,ワ
125EOS
126"#;
127
128        let cursor = Cursor::new(corpus_data.as_bytes());
129        let corpus = Corpus::from_reader(cursor).unwrap();
130
131        assert_eq!(corpus.len(), 2);
132        // Test basic properties without accessing private fields
133        // In a real implementation, we would add public accessor methods
134    }
135}
lindera_trainer/corpus.rs

lindera_trainer/
corpus.rs
