// 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.!
#include <cstdio>
#include <string>
#include "bpe_model.h"
#include "model_interface.h"
#include "testharness.h"
namespace sentencepiece {
namespace bpe {
namespace {
ModelProto MakeBaseModelProto() {
ModelProto model_proto;
auto *sp1 = model_proto.add_pieces();
auto *sp2 = model_proto.add_pieces();
auto *sp3 = model_proto.add_pieces();
sp1->set_type(ModelProto::SentencePiece::UNKNOWN);
sp1->set_piece("<unk>");
sp2->set_type(ModelProto::SentencePiece::CONTROL);
sp2->set_piece("<s>");
sp3->set_type(ModelProto::SentencePiece::CONTROL);
sp3->set_piece("</s>");
return model_proto;
}
void AddPiece(ModelProto *model_proto, const std::string &piece,
float score = 0.0) {
auto *sp = model_proto->add_pieces();
sp->set_piece(piece);
sp->set_score(score);
}
TEST(BPEModelTest, EncodeTest) {
ModelProto model_proto = MakeBaseModelProto();
AddPiece(&model_proto, "ab", 0.0); // 3
AddPiece(&model_proto, "cd", -0.1); // 4
AddPiece(&model_proto, "abc", -0.2); // 5
AddPiece(&model_proto, "a", -0.3); // 6
AddPiece(&model_proto, "b", -0.4); // 7
AddPiece(&model_proto, "c", -0.5); // 8
AddPiece(&model_proto, "ABC", -0.5); // 9
AddPiece(&model_proto, "abcdabcd", -0.5); // 10
AddPiece(&model_proto, "q", -0.5); // 11
AddPiece(&model_proto, "r", -0.5); // 12
AddPiece(&model_proto, "qr", -0.5); // 13
model_proto.mutable_pieces(9)->set_type( // ABC
ModelProto::SentencePiece::USER_DEFINED);
model_proto.mutable_pieces(10)->set_type( // abcdabcd
ModelProto::SentencePiece::USER_DEFINED);
model_proto.mutable_pieces(11)->set_type( // q
ModelProto::SentencePiece::USER_DEFINED);
model_proto.mutable_pieces(12)->set_type( // r
ModelProto::SentencePiece::USER_DEFINED);
const Model model(model_proto);
EncodeResult result;
result = model.Encode("");
EXPECT_TRUE(result.empty());
result = model.Encode("abc");
EXPECT_EQ(1, result.size());
EXPECT_EQ("abc", result[0].first);
result = model.Encode("AB");
EXPECT_EQ(2, result.size());
EXPECT_EQ("A", result[0].first);
EXPECT_EQ("B", result[1].first);
result = model.Encode("abcd");
EXPECT_EQ(2, result.size());
EXPECT_EQ("ab", result[0].first);
EXPECT_EQ("cd", result[1].first);
result = model.Encode("abcc");
EXPECT_EQ(2, result.size());
EXPECT_EQ("abc", result[0].first);
EXPECT_EQ("c", result[1].first);
result = model.Encode("xabcabaabcdd");
EXPECT_EQ(7, result.size());
EXPECT_EQ("x", result[0].first);
EXPECT_EQ("abc", result[1].first);
EXPECT_EQ("ab", result[2].first);
EXPECT_EQ("a", result[3].first);
EXPECT_EQ("ab", result[4].first);
EXPECT_EQ("cd", result[5].first);
EXPECT_EQ("d", result[6].first);
// all unknown.
result = model.Encode("xyz東京");
EXPECT_EQ(5, result.size());
EXPECT_EQ("x", result[0].first);
EXPECT_EQ("y", result[1].first);
EXPECT_EQ("z", result[2].first);
EXPECT_EQ("東", result[3].first);
EXPECT_EQ("京", result[4].first);
// User defined
result = model.Encode("ABC");
EXPECT_EQ(1, result.size());
EXPECT_EQ("ABC", result[0].first);
result = model.Encode("abABCcd");
EXPECT_EQ(3, result.size());
EXPECT_EQ("ab", result[0].first);
EXPECT_EQ("ABC", result[1].first);
EXPECT_EQ("cd", result[2].first);
// middle "abcdabcd" is user defined.
result = model.Encode("ababcdabcdcd");
EXPECT_EQ(3, result.size());
EXPECT_EQ("ab", result[0].first);
EXPECT_EQ("abcdabcd", result[1].first);
EXPECT_EQ("cd", result[2].first);
result = model.Encode("abqrcd");
EXPECT_EQ(4, result.size());
EXPECT_EQ("ab", result[0].first);
EXPECT_EQ("q", result[1].first);
EXPECT_EQ("r", result[2].first);
EXPECT_EQ("cd", result[3].first);
}
TEST(BPEModelTest, EncodeAmbiguousTest) {
ModelProto model_proto = MakeBaseModelProto();
AddPiece(&model_proto, "aa", -0.1);
AddPiece(&model_proto, "bb", -0.2);
AddPiece(&model_proto, "ab", -0.3);
AddPiece(&model_proto, "a", -0.4);
AddPiece(&model_proto, "b", -0.5);
const Model model(model_proto);
EncodeResult result;
// leftmost symbols are merged first.
result = model.Encode("aaa");
EXPECT_EQ(2, result.size());
EXPECT_EQ("aa", result[0].first);
EXPECT_EQ("a", result[1].first);
// "bb" is replaced earlier than "ab".
result = model.Encode("aabb");
EXPECT_EQ(2, result.size());
EXPECT_EQ("aa", result[0].first);
EXPECT_EQ("bb", result[1].first);
// "bb" is replaced earlier than "ab".
result = model.Encode("aaabbb");
EXPECT_EQ(4, result.size());
EXPECT_EQ("aa", result[0].first);
EXPECT_EQ("a", result[1].first);
EXPECT_EQ("bb", result[2].first);
EXPECT_EQ("b", result[3].first);
result = model.Encode("aaaba");
EXPECT_EQ(3, result.size());
EXPECT_EQ("aa", result[0].first);
EXPECT_EQ("ab", result[1].first);
EXPECT_EQ("a", result[2].first);
// makes a broken utf-8
const std::string broken_utf8 = std::string("あ").substr(0, 1);
result = model.Encode(broken_utf8);
EXPECT_EQ(1, result.size());
EXPECT_EQ(broken_utf8, result[0].first);
}
TEST(BPEModelTest, NotSupportedTest) {
ModelProto model_proto = MakeBaseModelProto();
const Model model(model_proto);
EXPECT_EQ(NBestEncodeResult(), model.NBestEncode("test", 10));
}
TEST(BPEModelTest, EncodeWithUnusedTest) {
ModelProto model_proto = MakeBaseModelProto();
AddPiece(&model_proto, "abcd", 10.0); // 3
AddPiece(&model_proto, "abc", 5.0); // 4
AddPiece(&model_proto, "ab", 2.0); // 5
AddPiece(&model_proto, "cd", 1.0); // 6
AddPiece(&model_proto, "a", 0.0); // 7
AddPiece(&model_proto, "b", 0.0); // 8
AddPiece(&model_proto, "c", 0.0); // 9
AddPiece(&model_proto, "d", 0.0); // 10
// No unused.
{
const Model model(model_proto);
const auto result = model.Encode("abcd");
EXPECT_EQ(1, result.size());
EXPECT_EQ("abcd", result[0].first);
}
{
model_proto.mutable_pieces(3)->set_type(ModelProto::SentencePiece::UNUSED);
const Model model(model_proto);
const auto result = model.Encode("abcd");
EXPECT_EQ(2, result.size());
EXPECT_EQ("abc", result[0].first);
EXPECT_EQ("d", result[1].first);
}
{
// The parent rule "abc" is still alive even if the child "ab" is unused.
model_proto.mutable_pieces(3)->set_type(ModelProto::SentencePiece::UNUSED);
model_proto.mutable_pieces(5)->set_type(ModelProto::SentencePiece::UNUSED);
const Model model(model_proto);
const auto result = model.Encode("abcd");
EXPECT_EQ(2, result.size());
EXPECT_EQ("abc", result[0].first);
EXPECT_EQ("d", result[1].first);
}
{
// This is tricky case. Even though "cd" is alive, it is not used, as
// it is not merged during the segmentation step.
// Segmentation: a|b|c|d => ab|c|d| => abc|d => abcd
// Resegmentation: abcd => abc|d => ab|c|d. ("abcd", "abc" are unsued)
model_proto.mutable_pieces(3)->set_type(ModelProto::SentencePiece::UNUSED);
model_proto.mutable_pieces(4)->set_type(ModelProto::SentencePiece::UNUSED);
model_proto.mutable_pieces(5)->set_type(ModelProto::SentencePiece::NORMAL);
const Model model(model_proto);
const auto result = model.Encode("abcd");
EXPECT_EQ(3, result.size());
EXPECT_EQ("ab", result[0].first);
EXPECT_EQ("c", result[1].first);
EXPECT_EQ("d", result[2].first);
}
}
TEST(SampleModelTest, EncodeTest) {
ModelProto model_proto = MakeBaseModelProto();
AddPiece(&model_proto, "ab", 0.0);
AddPiece(&model_proto, "cd", -0.1);
AddPiece(&model_proto, "abc", -0.2);
AddPiece(&model_proto, "abcd", -0.3);
// No regularization
{
const Model model(model_proto);
const auto result = model.Encode("abcd");
EXPECT_EQ(1, result.size());
EXPECT_EQ("abcd", result[0].first);
}
{
auto get_tokens = [](const EncodeResult& result) {
std::string out;
for (const auto& r : result) {
if (!result.empty()) {
out += ' ';
}
out += std::string(r.first);
}
return out;
};
const Model model(model_proto);
const std::vector<double> kAlpha = {0.0, 0.1, 0.5, 0.7, 0.9};
for (const auto alpha : kAlpha) {
constexpr int kTrial = 100000;
std::map<std::string, int> freq;
for (int n = 0; n < kTrial; ++n) {
freq[get_tokens(
model.SampleEncode("abcd", static_cast<float>(alpha)))]++;
}
int num = 0;
if (alpha == 0.0) {
EXPECT_EQ(1, freq.size());
} else {
EXPECT_GT(freq.size(), 1);
}
for (const auto& it : freq) {
num += it.second;
}
EXPECT_EQ(num, kTrial);
}
}
}
} // namespace
} // namespace bpe
} // namespace sentencepiece
Codebase Browser
chromium