// Copyright 2016 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.!

syntax = "proto2";

package sentencepiece;

// TrainerSpec encodes a various parameters for SentencePiece training.
message TrainerSpec {
  ///////////////////////////////////////////////////////////////////
  // General parameters
  //
  // Input corpus files.
  //  Trainer accepts the following two formats:
  //  A) Monolingual: plain text, one sentence per line.
  //  B) Bilingual:   TSV, source sentence <tab> target sentence
  //  When bilingual data is passed, shared vocabulary model is built.
  //  Note that the input file must be raw corpus, not a preprocessed corpus.
  //  Trainer only loads the first |input_sentence_size| sentences specified
  //  with this parameter.
  repeated string input  = 1;

  // Input corpus format:
  // "text": one-sentence-per-line text format (default)
  // "tsv":  sentence <tab> freq
  optional string input_format = 7;

  // Output model file prefix.
  // <model_prefix>.model and <model_prefix>.vocab are generated.
  optional string model_prefix = 2;

  // Model type. only have UNIGRAM now.
  enum ModelType {
    UNIGRAM = 1;   // Unigram language model with dynamic algorithm
    BPE     = 2;   // Byte Pair Encoding
    WORD    = 3;   // Delimitered by whitespace.
    CHAR    = 4;   // tokenizes into character sequence
  }
  optional ModelType model_type = 3 [ default = UNIGRAM ];

  // Vocabulary size. 8k is the default size.
  optional int32  vocab_size = 4 [ default = 8000 ];

  // List of the languages this model can accept.
  // Since the model is language-agnostic, this field is used as a reference.
  repeated string accept_language = 5;

  ///////////////////////////////////////////////////////////////////
  // Training parameters.
  //
  // Uses characters which cover the corpus with the ratio of |chars_coverage|.
  // This parameter determines the set of basic Alphabet of sentence piece.
  // 1.0 - |chars_coverage| characters are treated as UNK.
  optional float character_coverage = 10 [ default = 0.9995 ];

  // Maximum size of sentences the trainer loads from |input| parameter.
  // Trainer simply loads the |input| files in sequence.
  // It is better to shuffle the input corpus randomly.
  optional int32 input_sentence_size = 11 [ default = 10000000 ];

  // Maximum size of sentences to make seed sentence pieces.
  // Extended suffix array is constructed to extract frequent
  // sub-strings from the corpus. This uses 20N working space,
  // where N is the size of corpus.
  optional int32  mining_sentence_size = 12 [ default = 2000000 ];

  // Maximum size of sentences to train sentence pieces.
  optional int32  training_sentence_size = 13 [ default = 10000000 ];

  // The size of seed sentencepieces.
  // |seed_sentencepiece_size| must be larger than |vocab_size|.
  optional int32  seed_sentencepiece_size = 14 [ default = 1000000 ];

  // In every EM sub-iterations, keeps top
  // |shrinking_factor| * |current sentencepieces size| with respect to
  // the loss of the sentence piece. This value should be smaller than 1.0.
  optional float  shrinking_factor = 15 [ default = 0.75 ];

  // Number of threads in the training.
  optional int32  num_threads = 16 [ default = 16 ];

  // Number of EM sub iterations.
  optional int32  num_sub_iterations = 17 [ default = 2 ];

  ///////////////////////////////////////////////////////////////////
  // SentencePiece parameters which control the shapes of sentence piece.
  //
  // Maximum length of sentencepiece.
  optional int32  max_sentencepiece_length = 20 [ default = 16 ];

  // Uses Unicode script to split sentence pieces.
  // When |split_by_unicode_script| is true, we do not allow sentence piece to
  // include multiple Unicode scripts, e.g. "F1" is not a valid piece.
  // Exception: CJ characters (Hiragana/Katakana/Han) are all handled
  // as one script type, since Japanese word can consist of multiple scripts.
  // This exception is always applied regardless of the accept-language
  // parameter.
  optional bool split_by_unicode_script = 21 [ default = true ];

  // Use a white space to split sentence pieces.
  // When |split_by_whitespace| is false, we may have the piece containing
  // a white space in the middle. e.g., "in_the".
  optional bool split_by_whitespace = 22 [ default = true ];

  ///////////////////////////////////////////////////////////////////
  // Vocabulary management
  //
  // Defines control symbols used as an indicator to
  // change the behavior of the decoder. <s> and </s> are pre-defined.
  // We can use this field to encode various meta information,
  // including language indicator in multilingual model.
  // These symbols are not visible to users, but visible to
  // the decoder. Note that when the input sentence contains control symbols,
  // they are not treated as one token, but segmented into normal pieces.
  // Control symbols must be inserted independently from the segmentation.
  repeated string control_symbols = 30;

  // Defines user defined symbols.
  // These symbols are added with extremely high score
  // so they are always treated as one unique symbol in any context.
  // Typical usage of user_defined_symbols is placeholder for named entities.
  repeated string user_defined_symbols = 31;

  // `vocab_size` is treated as hard limit. Crash if
  // the model can not produce the vocab of size `vocab_size`,
  // When `hard_vocab_limit` is false, vocab_size is treated
  // as soft limit. Note that when model_type=char,
  // always assumes hard_vocab_limit = false.
  optional bool hard_vocab_limit = 33 [ default = true ];

  ///////////////////////////////////////////////////////////////////
  // Reserved special meta tokens.
  // * -1 is not used.
  // * unk_id must not be -1.
  // Id must starts with 0 and be contigous.
  optional int32 unk_id = 40 [ default = 0 ];    // <unk>
  optional int32 bos_id = 41 [ default = 1 ];    // <s>
  optional int32 eos_id = 42 [ default = 2 ];    // </s>
  optional int32 pad_id = 43 [ default = -1 ];   // <pad> (padding)

  // Encodes <unk> into U+2047 (DOUBLE QUESTION MARK),
  // since this character can be useful both for user and
  // developer. We can easily figure out that <unk> is emitted.
  optional string unk_surface = 44 [ default = " \xE2\x81\x87 " ];

  // Customized extensions: the range of field numbers
  // are open to third-party extensions.
  extensions 200 to max;
};

// NormalizerSpec encodes a various parameters for string normalizaiton
message NormalizerSpec {
  // name of normalization rule.
  optional string name = 1;

  // Pre-compiled normalization rule created by
  // Builder::GetPrecompiledCharsMap() or Builder::CompileCharsMap() method.
  // Usually this field is set by Builder::GetNormalizerSpec() method.
  optional bytes precompiled_charsmap = 2;

  // Adds dummy whitespace at the beginning of text in order to
  // treat "world" in "world" and "hello world" in the same way.
  optional bool add_dummy_prefix = 3 [ default = true ];

  // Removes leading, trailing, and duplicate internal whitespace.
  optional bool remove_extra_whitespaces = 4  [ default = true ];

  // Replaces whitespace with meta symbol.
  // This field must be true to train sentence piece model.
  optional bool escape_whitespaces = 5 [ default = true ];

  // Custom normalization rule file in TSV format.
  // https://github.com/google/sentencepiece/blob/master/doc/normalization.md
  // This field is only used in SentencePieceTrainer::Train() method, which
  // compiles the rule into the binary rule stored in `precompiled_charsmap`.
  optional string normalization_rule_tsv = 6;

  // Customized extensions: the range of field numbers
  // are open to third-party extensions.
  extensions 200 to max;
}

// ModelProto stores model parameters.
// SentencePieceProcessor is supposed to be self-contained.
// All settings/parameters which may change the behavior must be encoded
// in ModelProto.
message ModelProto {
  message SentencePiece {
    enum Type {
      NORMAL       = 1;  // normal symbol
      UNKNOWN      = 2;  // unknown symbol. only <unk> for now.
      CONTROL      = 3;  // control symbols. </s>, <s>, <2ja> etc.
      USER_DEFINED = 4;  // user defined symbols.
                         // Typical usage of USER_DEFINED symbol
                         // is placeholder.
      UNUSED       = 5;  // this piece is not used.
    };
    optional string piece = 1;   // piece must not be empty.
    optional float  score = 2;
    optional Type   type = 3 [ default =  NORMAL ];

    // Customized extensions: the range of field numbers
    // are open to third-party extensions.
    extensions 200 to max;
  }

  // Sentence pieces with scores.
  repeated SentencePiece pieces = 1;

  // Spec used to generate this model file.
  optional TrainerSpec trainer_spec = 2;

  // Spec for text normalization.
  optional NormalizerSpec normalizer_spec = 3;

  // Customized extensions: the range of field numbers
  // are open to third-party extensions.
  extensions 200 to max;
}