EOS_ID and UNK_ID are now determined from the loaded vocabulary file, instead of requiring hard-coded values

1 files changed, 6 insertions, 0 deletions

diff --git a/src/data/vocab.h b/src/data/vocab.h
index 89dc9ccb..391083a1 100644
--- a/src/data/vocab.h
+++ b/src/data/vocab.h
@@ -35,6 +35,9 @@ public:
   void create(InputFileStream& trainStrm,
               OutputFileStream& vocabStrm,
               size_t maxSize = 0);
+    
+  Word GetEosId() const { return eosId_; }
+  Word GetUnkId() const { return unkId_; }
 
 private:
   typedef std::map<std::string, size_t> Str2Id;
@@ -43,6 +46,9 @@ private:
   typedef std::vector<std::string> Id2Str;
   Id2Str id2str_;
 
+  Word eosId_ = -1;
+  Word unkId_ = -1;
+
   class VocabFreqOrderer;
 };
 }