Add tag_dictionary argument to Trainer::new()

README-ja.md

@@ -21,40 +21,40 @@ Vaporetto はトークン化モデルを生成するための方法を3つ用意
 
 #### 配布モデルをダウンロードする
 
-1番目は最も単純な方法で、我々によって学習されたモデルをダウンロードすることです。
+1つ目は最も単純な方法で、学習済みモデルをダウンロードすることです。
 モデルファイルは[ここ](https://github.com/daac-tools/vaporetto/releases)にあります。
 
 `bccwj-suw+unidic+tag` を選びました。
 ```
 % wget https://github.com/daac-tools/vaporetto/releases/download/v0.5.0/bccwj-suw+unidic+tag.tar.xz
 ```
 
-各ファイルにはモデルファイルとライセンス条項が含まれているので、以下のようなコマンドでダウンロードしたファイルを展開する必要があります。
+各ファイルはモデルファイルとライセンス条項が含まれた圧縮ファイルなので、ダウンロードしたファイルを展開する必要があります。
 ```
 % tar xf ./bccwj-suw+unidic+tag.tar.xz
 ```
 
-トークン化を行うには、以下のコマンドを実行します。
+トークン化には、以下のコマンドを実行します。
 ```
 % echo 'ヴェネツィアはイタリアにあります。' | cargo run --release -p predict -- --model path/to/bccwj-suw+unidic+tag.model.zst
 ```
 
-以下が出力されるでしょう。
+以下が出力されます。
 ```
 ヴェネツィア は イタリア に あり ます 。
 ```
 
 #### KyTea のモデルを変換する
 
-2番目の方法も単純で、 KyTea で学習されたモデルを変換することです。
+2つ目の方法も単純で、 KyTea で学習されたモデルを変換することです。
 まずはじめに、好きなモデルを [KyTea Models](http://www.phontron.com/kytea/model.html) ページからダウンロードします。
 
 `jp-0.4.7-5.mod.gz` を選びました。
 ```
 % wget http://www.phontron.com/kytea/download/model/jp-0.4.7-5.mod.gz
 ```
 
-各モデルは圧縮されているので、以下のようなコマンドでダウンロードしたモデルを展開する必要があります。
+各モデルは圧縮されているので、ダウンロードしたモデルを展開する必要があります。
 ```
 % gunzip ./jp-0.4.7-5.mod.gz
 ```
@@ -64,21 +64,21 @@ KyTea のモデルを Vaporetto のモデルに変換するには、 Vaporetto 
 % cargo run --release -p convert_kytea_model -- --model-in path/to/jp-0.4.7-5.mod --model-out path/to/jp-0.4.7-5-tokenize.model.zst
 ```
 
-これでトークン化を行えます。以下のコマンドを実行します。
+これでトークン化できます。以下のコマンドを実行します。
 ```
 % echo 'ヴェネツィアはイタリアにあります。' | cargo run --release -p predict -- --model path/to/jp-0.4.7-5-tokenize.model.zst
 ```
 
-以下が出力されるでしょう。
+以下が出力されます。
 ```
 ヴェネツィア は イタリア に あ り ま す 。
 ```
 
 #### 自分のモデルを学習する
 
-3番目は主に研究者向けで、自分で学習コーパスを用意し、自分でトークン化モデルを学習することです。
+3つ目は主に研究者向けで、自分で学習コーパスを用意し、モデルを学習することです。
 
-Vaporetto は2種類のコーパス、すなわちフルアノテーションコーパスと部分アノテーションコーパスから学習することが可能です。
+Vaporetto は2種類のコーパス（フルアノテーションコーパスと部分アノテーションコーパス）から学習することが可能です。
 
 フルアノテーションコーパスは、すべての文字境界に対してトークン境界であるかトークンの内部であるかがアノテーションされたコーパスです。
 このデータは、以下に示すようにトークン境界に空白が挿入された形式です。
@@ -88,7 +88,7 @@ Vaporetto は2種類のコーパス、すなわちフルアノテーションコ
 火星 猫 の 生態 の 調査 結果
 ```
 
-一方、部分アノテーションコーパスは一部の文字境界のみに対してアノテーションされたコーパスです。
+部分アノテーションコーパスは、一部の文字境界のみに対してアノテーションされたコーパスです。
 各文字境界には `|` (トークン境界)、 `-` (非トークン境界)、 ` ` (不明) のいずれかの形式でアノテーションされます。
 
 ここに例を示します。
@@ -104,7 +104,13 @@ Vaporetto は2種類のコーパス、すなわちフルアノテーションコ
 
 `--tok` 引数ではフルアノテーションコーパスを指定し、 `--part` 引数では部分アノテーションコーパスを指定します。
 `--dict` 引数によって単語辞書を指定することもできます。
-単語辞書は、1行1単語のファイルです。
+単語辞書は、1行1単語のファイルであり、必要に応じてタグを付与することもできます。
+```
+トスカーナ
+パンツァーノ
+灯里/名詞-固有名詞-人名-名/アカリ
+形態/名詞-普通名詞-一般/ケータイ
+```
 
 学習器は空行の入力を受け付けません。
 このため、学習の前にコーパスから空行を削除してください。

README.md

@@ -105,7 +105,14 @@ To train a model, use the following command:
 
 The `--tok` argument specifies a fully annotated corpus, and the `--part` argument specifies a partially annotated corpus.
 You can also specify a word dictionary with the `--dict` argument.
-A word dictionary is a file with words per line.
+A word dictionary is a file that lists words line by line and can be tagged as needed:
+
+```
+トスカーナ
+パンツァーノ
+灯里/名詞-固有名詞-人名-名/アカリ
+形態/名詞-普通名詞-一般/ケータイ
+```
 
 The trainer does not accept empty lines.
 Therefore, remove all empty lines from the corpus before training.

examples/wasm/src/lib.rs

@@ -12,23 +12,21 @@ use yew::{html, Component, Context, Html};
 use once_cell::sync::Lazy;
 use vaporetto::{CharacterType, Model, Predictor, Sentence};
 use vaporetto_rules::{
-    sentence_filters::{ConcatGraphemeClustersFilter, KyteaWsConstFilter, PatternMatchTagger},
+    sentence_filters::{ConcatGraphemeClustersFilter, KyteaWsConstFilter},
     string_filters::KyteaFullwidthFilter,
     SentenceFilter, StringFilter,
 };
 
 use crate::text_input::TextInput;
 use crate::token_view::TokenView;
 
-static PREDICTOR: Lazy<(Predictor, PatternMatchTagger)> = Lazy::new(|| {
+static PREDICTOR: Lazy<Predictor> = Lazy::new(|| {
     let mut f = Cursor::new(include_bytes!("bccwj-suw+unidic+tag-huge.model.zst"));
     let mut decoder = ruzstd::StreamingDecoder::new(&mut f).unwrap();
     let mut buff = vec![];
     decoder.read_to_end(&mut buff).unwrap();
-    let (model, rest) = Model::read_slice(&buff).unwrap();
-    let config = bincode::config::standard();
-    let word_tag_map: Vec<(String, Vec<Option<String>>)> = bincode::decode_from_slice(rest, config).unwrap().0;
-    (Predictor::new(model, true).unwrap(), PatternMatchTagger::new(word_tag_map.into_iter().collect()))
+    let (model, _) = Model::read_slice(&buff).unwrap();
+    Predictor::new(model, true).unwrap()
 });
 
 pub enum Message {
@@ -84,15 +82,14 @@ impl gloo_worker::Worker for Worker {
         let filtered_text = pre_filter.filter(sentence_orig.as_raw_text());
         sentence_filtered.update_raw(filtered_text).unwrap();
 
-        PREDICTOR.0.predict(sentence_filtered);
+        PREDICTOR.predict(sentence_filtered);
 
         let wsconst_g = ConcatGraphemeClustersFilter;
         let wsconst_d = KyteaWsConstFilter::new(CharacterType::Digit);
         wsconst_g.filter(sentence_filtered);
         wsconst_d.filter(sentence_filtered);
 
         sentence_filtered.fill_tags();
-        PREDICTOR.1.filter(sentence_filtered);
         let n_tags = sentence_filtered.n_tags();
 
         sentence_orig

predict/src/main.rs

@@ -7,7 +7,7 @@ use std::time::Instant;
 use clap::Parser;
 use vaporetto::{CharacterType, Model, Predictor, Sentence};
 use vaporetto_rules::{
-    sentence_filters::{ConcatGraphemeClustersFilter, KyteaWsConstFilter, PatternMatchTagger},
+    sentence_filters::{ConcatGraphemeClustersFilter, KyteaWsConstFilter},
     string_filters::KyteaFullwidthFilter,
     SentenceFilter, StringFilter,
 };
@@ -92,14 +92,6 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     let mut f = zstd::Decoder::new(File::open(args.model)?)?;
     let model = Model::read(&mut f)?;
     let predictor = Predictor::new(model, args.predict_tags)?;
-    let word_tag_map: Vec<(String, Vec<Option<String>>)> = if args.predict_tags {
-        let config = bincode::config::standard();
-        bincode::decode_from_std_read(&mut f, config).unwrap_or_else(|_| vec![])
-    } else {
-        vec![]
-    };
-    let pattern_match_tagger = (!word_tag_map.is_empty())
-        .then(|| PatternMatchTagger::new(word_tag_map.into_iter().collect()));
 
     eprintln!("Start tokenization");
     let start = Instant::now();
@@ -119,9 +111,6 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                 post_filters.iter().for_each(|filter| filter.filter(&mut s));
                 if args.predict_tags {
                     s.fill_tags();
-                    if let Some(tagger) = pattern_match_tagger.as_ref() {
-                        tagger.filter(&mut s);
-                    }
                 }
                 s.write_tokenized_text(&mut buf);
                 writeln!(out, "{}", buf)?;
@@ -142,9 +131,6 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                 post_filters.iter().for_each(|filter| filter.filter(&mut s));
                 if args.predict_tags {
                     s.fill_tags();
-                    if let Some(tagger) = pattern_match_tagger.as_ref() {
-                        tagger.filter(&mut s);
-                    }
                 }
                 s_orig.update_raw(line)?;
                 s_orig.reset_tags(s.n_tags());

train/src/main.rs

@@ -1,4 +1,4 @@
-use std::collections::BTreeMap;
+use std::collections::BTreeSet;
 use std::fs::File;
 use std::io::{prelude::*, stderr, BufReader};
 use std::path::PathBuf;
@@ -125,42 +125,46 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         eprintln!("# of sentences: {}", train_sents.len());
     }
 
-    let mut word_tag_map = BTreeMap::new();
-    let mut word_buf = Sentence::default();
+    let mut tag_dictionary = vec![];
+    let mut dictionary = BTreeSet::new();
     for path in args.dict {
         eprintln!("Loading {:?} ...", path);
         let f = File::open(path)?;
         let f = BufReader::new(f);
-        for (i, line) in f.lines().enumerate() {
-            if i % 100000 == 0 {
-                eprint!("# of words: {}\r", i);
+        for line in f.lines() {
+            if dictionary.len() % 10000 == 0 {
+                eprint!("# of words: {}\r", dictionary.len());
                 stderr().flush()?;
             }
-            let line = line?;
-            word_buf.update_tokenized(&line).unwrap();
-            for token in word_buf.iter_tokens() {
-                let word = token.surface().to_string();
-                let word = if args.no_norm {
-                    word
-                } else {
-                    fullwidth_filter.filter(&word)
-                };
-                word_tag_map.entry(word).or_insert_with(|| {
-                    token
-                        .tags()
-                        .iter()
-                        .map(|tag| tag.as_ref().map(|tag| tag.to_string()))
-                        .collect()
-                });
+            let s = Sentence::from_tokenized(&line?)?;
+            let s = if args.no_norm {
+                s
+            } else {
+                let new_line = fullwidth_filter.filter(s.as_raw_text());
+                let mut new_s = Sentence::from_raw(new_line)?;
+                new_s.boundaries_mut().clone_from_slice(s.boundaries());
+                new_s.reset_tags(s.n_tags());
+                new_s.tags_mut().clone_from_slice(s.tags());
+                new_s
+            };
+            for token in s.iter_tokens() {
+                dictionary.insert(token.surface().to_string());
             }
+            tag_dictionary.push(s);
         }
-        eprintln!("# of words: {}", word_tag_map.len());
+        eprintln!("# of words: {}", dictionary.len());
     }
-    let dictionary = word_tag_map.iter().map(|(word, _)| word.clone()).collect();
+    let dictionary = dictionary.into_iter().collect();
 
     eprintln!("Extracting into features...");
     let mut trainer = Trainer::new(
-        args.charw, args.charn, args.typew, args.typen, dictionary, args.dictn,
+        args.charw,
+        args.charn,
+        args.typew,
+        args.typen,
+        dictionary,
+        args.dictn,
+        &tag_dictionary,
     )?;
     for (i, s) in train_sents.iter().enumerate() {
         if i % 10000 == 0 {
@@ -177,12 +181,6 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let mut f = zstd::Encoder::new(File::create(args.model)?, 19)?;
     model.write(&mut f)?;
-    for tag_model in model.tag_models() {
-        word_tag_map.remove(tag_model.token());
-    }
-    let word_tag_map: Vec<(String, Vec<Option<String>>)> = word_tag_map.into_iter().collect();
-    let config = bincode::config::standard();
-    bincode::encode_into_std_write(&word_tag_map, &mut f, config)?;
     f.finish()?;
 
     Ok(())

vaporetto/src/tag_trainer.rs

@@ -47,6 +47,7 @@ pub struct TagTrainer<'a> {
     _type_window_size: u8,
     type_ngram_size: u8,
     // Uses BTreeMap to improve compression ratio.
+    default_tags: HashMap<&'a str, &'a [Option<Cow<'a, str>>]>,
     examples: BTreeMap<&'a str, Vec<TagExample<'a>>>,
 }
 
@@ -56,12 +57,14 @@ impl<'a> TagTrainer<'a> {
         char_ngram_size: u8,
         type_window_size: u8,
         type_ngram_size: u8,
+        default_tags: HashMap<&'a str, &'a [Option<Cow<'a, str>>]>,
     ) -> Self {
         Self {
             _char_window_size: char_window_size,
             char_ngram_size,
             _type_window_size: type_window_size,
             type_ngram_size,
+            default_tags,
             examples: BTreeMap::new(),
         }
     }
@@ -295,7 +298,18 @@ impl<'a> TagTrainer<'a> {
         })
     }
 
-    pub fn train(self, epsilon: f64, cost: f64, solver: SolverType) -> Result<Vec<TagModel>> {
+    pub fn train(mut self, epsilon: f64, cost: f64, solver: SolverType) -> Result<Vec<TagModel>> {
+        for (token, tags) in self.default_tags {
+            if tags.iter().any(|t| t.is_some()) && !self.examples.contains_key(token) {
+                self.examples.insert(
+                    token,
+                    vec![TagExample {
+                        tags,
+                        features: vec![],
+                    }],
+                );
+            }
+        }
         let mut tag_models = vec![];
         liblinear::toggle_liblinear_stdout_output(false);
         let n_tokens = self.examples.len();