add function to load vocabulary

Author: mrityunjay-tripathi

models/bert/CMakeLists.txt

@@ -8,6 +8,7 @@ set(SOURCES
   bert.hpp
   bert_impl.hpp
   bert_tokenizer.hpp
+  bert_tokenizer_impl.hpp
 )
 
 foreach(file ${SOURCES})

models/bert/bert.hpp

@@ -23,17 +23,12 @@ namespace ann /** Artificial Neural Network. */ {
 
 /**
  * @tparam OutputLayerType Type of the last layer to be added to BERT model.
- * @tparam InitType Initilization Rule to be used to initialize parameters.
- * @tparam InputDataType Type of the input data (arma::colvec, arma::mat,
- *         arma::sp_mat or arma::cube).
- * @tparam OutputDataType Type of the output data (arma::colvec, arma::mat,
- *         arma::sp_mat or arma::cube).
+ * @tparam InitializationRuleType Initilization Rule to be used to initialize
+ *         parameters.
  */
 template <
   typename OutputLayerType = NegativeLogLikelihood<>,
-  typename InitializationRuleType = XavierInitialization,
-  typename InputDataType = arma::mat,
-  typename OutputDataType = arma::mat
+  typename InitializationRuleType = XavierInitialization
 >
 class BERT
 {
@@ -58,8 +53,8 @@ class BERT
        const size_t dModel = 512,
        const size_t numHeads = 8,
        const double dropout = 0.1,
-       const InputDataType& attentionMask = InputDataType(),
-       const InputDataType& keyPaddingMask = InputDataType());
+       const arma::mat& attentionMask = arma::mat(),
+       const arma::mat& keyPaddingMask = arma::mat());
 
   /**
    * Load the network from a local directory.
@@ -98,10 +93,10 @@ class BERT
   double dropout;
 
   //! Locally-stored attention mask.
-  InputDataType attentionMask;
+  arma::mat attentionMask;
 
   //! Locally-stored key padding mask.
-  InputDataType keyPaddingMask;
+  arma::mat keyPaddingMask;
 
   //! Locally-stored complete decoder network.
   FFN<OutputLayerType, InitializationRuleType> bert;

models/bert/bert_impl.hpp

@@ -19,9 +19,8 @@
 namespace mlpack {
 namespace ann /** Artificial Neural Network. */ {
 
-template<typename OutputLayerType, typename InitType, typename InputDataType,
-    typename OutputDataType>
-BERT<OutputLayerType, InitType, InputDataType, OutputDataType>::BERT() :
+template<typename OutputLayerType, typename InitializationRuleType>
+BERT<OutputLayerType, InitializationRuleType>::BERT() :
     srcVocabSize(0),
     srcSeqLen(0),
     numEncoderLayers(0),
@@ -33,17 +32,16 @@ BERT<OutputLayerType, InitType, InputDataType, OutputDataType>::BERT() :
   // Nothing to do here.
 }
 
-template<typename OutputLayerType, typename InitType, typename InputDataType,
-    typename OutputDataType>
-BERT<OutputLayerType, InitType, InputDataType, OutputDataType>::BERT(
+template<typename OutputLayerType, typename InitializationRuleType>
+BERT<OutputLayerType, InitializationRuleType>::BERT(
     const size_t srcVocabSize,
     const size_t srcSeqLen,
     const size_t numEncoderLayers,
     const size_t dModel,
     const size_t numHeads,
     const double dropout,
-    const InputDataType& attentionMask,
-    const InputDataType& keyPaddingMask) :
+    const arma::mat& attentionMask,
+    const arma::mat& keyPaddingMask) :
     srcVocabSize(srcVocabSize),
     srcSeqLen(srcSeqLen),
     numEncoderLayers(numEncoderLayers),
@@ -72,25 +70,22 @@ BERT<OutputLayerType, InitType, InputDataType, OutputDataType>::BERT(
       dimFFN,
       dropout,
       attentionMask,
-      keyPaddingMask
-    );
+      keyPaddingMask);
 
     bert.Add(encoder.Model());
   }
 }
 
-template<typename OutputLayerType, typename InitType, typename InputDataType,
-    typename OutputDataType>
-void BERT<OutputLayerType, InitType, InputDataType, OutputDataType>::LoadModel(
+template<typename OutputLayerType, typename InitializationRuleType>
+void BERT<OutputLayerType, InitializationRuleType>::LoadModel(
     const std::string& filepath)
 {
   data::Load(filepath, "BERT", bert);
   std::cout << "Loaded model" << std::endl;
 }
 
-template<typename OutputLayerType, typename InitType, typename InputDataType,
-    typename OutputDataType>
-void BERT<OutputLayerType, InitType, InputDataType, OutputDataType>::SaveModel(
+template<typename OutputLayerType, typename InitializationRuleType>
+void BERT<OutputLayerType, InitializationRuleType>::SaveModel(
     const std::string& filepath)
 {
   std::cout << "Saving model" << std::endl;

models/bert/bert_tokenizer.hpp

@@ -44,6 +44,9 @@ template <
 class BertTokenizer
 {
  public:
+  /**
+   * Create a BertTokenizer object.
+   */
   BertTokenizer();
 
   /**
@@ -124,6 +127,8 @@ class BertTokenizer
   std::string maskToken;
 }; // class BertTokenizer
 
+// Include implementation.
+#include "bert_tokenizer_impl.hpp"
 } // namespace ann
 } // namespace mlpack
 

models/bert/bert_tokenizer_impl.hpp

@@ -0,0 +1,62 @@
+/**
+ * @file models/bert/bert_tokenizer_impl.hpp
+ * @author Mrityunjay Tripathi
+ *
+ * Implementation of the BERT Tokenizer.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+
+#ifndef MODELS_BERT_BERT_TOKENIZER_IMPL_HPP
+#define MODELS_BERT_BERT_TOKENIZER_IMPL_HPP
+
+#include "bert_tokenizer.hpp"
+
+namespace mlpack {
+namespace ann /** Artificial Neural Network. */ {
+
+template<typename InputDataType, typename OutputDataType>
+BertTokenizer<InputDataType, OutputDataType>::BertTokenizer() :
+    vocabFile(""),
+    lowerCase(true),
+    basicTokenize(true),
+    unkToken("[UNK]"),
+    sepToken("[SEP]"),
+    padToken("[PAD]"),
+    clsToken("[CLS]"),
+    maskToken("[MASK]")
+{
+  // Nothing to do here.
+}
+
+template<typename InputDataType, typename OutputDataType>
+BertTokenizer<InputDataType, OutputDataType>::BertTokenizer(
+    const std::string vocabFile,
+    const bool lowerCase,
+    const bool basicTokenize,
+    const std::vector<std::string> neverSplit,
+    const std::string unkToken,
+    const std::string sepToken,
+    const std::string padToken,
+    const std::string clsToken,
+    const std::string maskToken) :
+    vocabFile(vocabFile),
+    lowerCase(lowerCase),
+    basicTokenize(basicTokenize),
+    neverSplit(neverSplit),
+    unkToken(unkToken),
+    sepToken(sepToken),
+    padToken(padToken),
+    clsToken(clsToken),
+    maskToken(maskToken)
+{
+  // code here.
+}
+
+} // namespace ann
+} // namespace mlpack
+
+#endif

utils/utils.hpp

@@ -29,7 +29,7 @@ class Utils
  public:
   /**
    * Determines whether a path exists.
-   * 
+   *
    * @param path Global or relative path.
    * @param absolutePath Boolean to determine if path is absolute or relative.
    * @return true if path exists else false. Defaults to false.
@@ -329,5 +329,28 @@ class Utils
       mlpack::Log::Warn << "The " << path << " doesn't exist." << std::endl;
     }
   }
+
+  /**
+   * Loads a vocabulary file and stores the content into a vector.
+   *
+   * @param vocabPath The path to the vocabulary file.
+   * @param vocabulary Stores the vocabulary content along with indices.
+   */
+  void LoadVocabulary(const std::string vocabPath,
+                      std::map<std::string, size_t>& vocabulary)
+  {
+    std::string token;
+    std::ifstream vocabFile(vocabPath);
+    if (vocabFile.is_open())
+    {
+      for (size_t i = 0; std::getline(vocabFile, token); ++i)
+        vocabulary[token] = i;
+
+      vocabFile.close();
+    }
+
+    else
+      std::cout << "Unable to open vocabulary file!" << std::endl;
+  }
 };
 #endif