Support to change ids of <unk>, <s>, </s>

1 files changed, 1 insertions, 1 deletions

diff --git a/src/bpe_model.cc b/src/bpe_model.cc
index fdc5fd7..fe9df73 100644
--- a/src/bpe_model.cc
+++ b/src/bpe_model.cc
@@ -22,7 +22,6 @@ namespace bpe {
 
 Model::Model(const ModelProto &model_proto) {
   model_proto_ = &model_proto;
-  CheckControlSymbols();
 
   for (int i = 0; i < model_proto_->pieces_size(); ++i) {
     const auto &sp = model_proto_->pieces(i);
@@ -35,6 +34,7 @@ Model::Model(const ModelProto &model_proto) {
       LOG(FATAL) << "User defined symbol is not supported in BPE";
     } else {
       port::InsertOrDie(&reserved_id_map_, sp.piece(), i);
+      if (sp.type() == ModelProto::SentencePiece::UNKNOWN) unk_id_ = i;
     }
   }
 }