Add option for n-gram length

1 files changed, 6 insertions, 3 deletions

diff --git a/lm/builder/dump_counts_main.cc b/lm/builder/dump_counts_main.cc
index 6c0b3a4..200b66a 100644
--- a/lm/builder/dump_counts_main.cc
+++ b/lm/builder/dump_counts_main.cc
@@ -2,14 +2,17 @@
 #include "util/file.hh"
 #include "util/stream/chain.hh"
 #include "util/stream/io.hh"
+
 #include <iostream>
 
+#include <boost/lexical_cast.hpp>
+
 int main(int argc, char *argv[]) {
-  if (argc != 2) {
-    std::cerr << "Expected null-delimited vocab file." << std::endl;
+  if (argc != 3) {
+    std::cerr << "Expected null-delimited vocab file and order." << std::endl;
     return 1;
   }
   util::scoped_fd vocab(util::OpenReadOrThrow(argv[1]));
   lm::builder::VocabReconstitute reconstitute(vocab.get());
-  util::stream::Chain(util::stream::ChainConfig(5*4+8, 1, 4096)) >> util::stream::Read(0) >> lm::builder::Print<uint64_t>(reconstitute, std::cout);
+  util::stream::Chain(util::stream::ChainConfig(boost::lexical_cast<unsigned int>(argv[2])*4+8, 1, 4096)) >> util::stream::Read(0) >> lm::builder::Print<uint64_t>(reconstitute, std::cout);
 }