diff options
Diffstat (limited to 'src/trainer_interface_test.cc')
-rw-r--r-- | src/trainer_interface_test.cc | 30 |
1 files changed, 30 insertions, 0 deletions
diff --git a/src/trainer_interface_test.cc b/src/trainer_interface_test.cc index 34feb85..45d4905 100644 --- a/src/trainer_interface_test.cc +++ b/src/trainer_interface_test.cc @@ -288,6 +288,36 @@ TEST(TrainerInterfaceTest, OverrideSpecialPiecesTest) { EXPECT_EQ(ModelProto::SentencePiece::USER_DEFINED, trainer.meta_pieces_[1].second); } + + { + auto trainer_spec = base_trainer_spec; + trainer_spec.set_unk_piece("__UNK__"); + trainer_spec.set_bos_piece("__BOS__"); + trainer_spec.set_eos_piece("__EOS__"); + trainer_spec.set_pad_piece("__PAD__"); + trainer_spec.set_pad_id(3); + TrainerInterface trainer(trainer_spec, normalizer_spec); + EXPECT_TRUE(trainer.status().ok()); + EXPECT_EQ("__UNK__", trainer.meta_pieces_[0].first); + EXPECT_EQ("__BOS__", trainer.meta_pieces_[1].first); + EXPECT_EQ("__EOS__", trainer.meta_pieces_[2].first); + EXPECT_EQ("__PAD__", trainer.meta_pieces_[3].first); + } + + { + auto trainer_spec = base_trainer_spec; + trainer_spec.set_unk_piece("__UNK__"); + trainer_spec.set_bos_piece("__UNK__"); + TrainerInterface trainer(trainer_spec, normalizer_spec); + EXPECT_FALSE(trainer.status().ok()); + } + + { + auto trainer_spec = base_trainer_spec; + trainer_spec.set_unk_piece(""); + TrainerInterface trainer(trainer_spec, normalizer_spec); + EXPECT_FALSE(trainer.status().ok()); + } } TEST(TrainerInterfaceTest, SerializeTest) { |