1 files changed, 37 insertions, 14 deletions

diff --git a/klm/lm/query_main.cc b/klm/lm/query_main.cc
index 27d3a1a5..bd4fde62 100644
--- a/klm/lm/query_main.cc
+++ b/klm/lm/query_main.cc
@@ -1,42 +1,65 @@
 #include "lm/ngram_query.hh"
 
+#ifdef WITH_NPLM
+#include "lm/wrappers/nplm.hh"
+#endif
+
+#include <stdlib.h>
+
+void Usage(const char *name) {
+  std::cerr << "KenLM was compiled with maximum order " << KENLM_MAX_ORDER << "." << std::endl;
+  std::cerr << "Usage: " << name << " [-n] lm_file" << std::endl;
+  std::cerr << "Input is wrapped in <s> and </s> unless -n is passed." << std::endl;
+  exit(1);
+}
+
 int main(int argc, char *argv[]) {
-  if (!(argc == 2 || (argc == 3 && !strcmp(argv[2], "null")))) {
-    std::cerr << "KenLM was compiled with maximum order " << KENLM_MAX_ORDER << "." << std::endl;
-    std::cerr << "Usage: " << argv[0] << " lm_file [null]" << std::endl;
-    std::cerr << "Input is wrapped in <s> and </s> unless null is passed." << std::endl;
-    return 1;
+  bool sentence_context = true;
+  const char *file = NULL;
+  for (char **arg = argv + 1; arg != argv + argc; ++arg) {
+    if (!strcmp(*arg, "-n")) {
+      sentence_context = false;
+    } else if (!strcmp(*arg, "-h") || !strcmp(*arg, "--help") || file) {
+      Usage(argv[0]);
+    } else {
+      file = *arg;
+    }
   }
+  if (!file) Usage(argv[0]);
   try {
-    bool sentence_context = (argc == 2);
     using namespace lm::ngram;
     ModelType model_type;
-    if (RecognizeBinary(argv[1], model_type)) {
+    if (RecognizeBinary(file, model_type)) {
       switch(model_type) {
         case PROBING:
-          Query<lm::ngram::ProbingModel>(argv[1], sentence_context, std::cin, std::cout);
+          Query<lm::ngram::ProbingModel>(file, sentence_context, std::cin, std::cout);
           break;
         case REST_PROBING:
-          Query<lm::ngram::RestProbingModel>(argv[1], sentence_context, std::cin, std::cout);
+          Query<lm::ngram::RestProbingModel>(file, sentence_context, std::cin, std::cout);
           break;
         case TRIE:
-          Query<TrieModel>(argv[1], sentence_context, std::cin, std::cout);
+          Query<TrieModel>(file, sentence_context, std::cin, std::cout);
           break;
         case QUANT_TRIE:
-          Query<QuantTrieModel>(argv[1], sentence_context, std::cin, std::cout);
+          Query<QuantTrieModel>(file, sentence_context, std::cin, std::cout);
           break;
         case ARRAY_TRIE:
-          Query<ArrayTrieModel>(argv[1], sentence_context, std::cin, std::cout);
+          Query<ArrayTrieModel>(file, sentence_context, std::cin, std::cout);
           break;
         case QUANT_ARRAY_TRIE:
-          Query<QuantArrayTrieModel>(argv[1], sentence_context, std::cin, std::cout);
+          Query<QuantArrayTrieModel>(file, sentence_context, std::cin, std::cout);
           break;
         default:
           std::cerr << "Unrecognized kenlm model type " << model_type << std::endl;
           abort();
       }
+#ifdef WITH_NPLM
+    } else if (lm::np::Model::Recognize(file)) {
+      lm::np::Model model(file);
+      Query(model, sentence_context, std::cin, std::cout);
+#endif
     } else {
-      Query<ProbingModel>(argv[1], sentence_context, std::cin, std::cout);
+      Query<ProbingModel>(file, sentence_context, std::cin, std::cout);
     }
     std::cerr << "Total time including destruction:\n";
     util::PrintUsage(std::cerr);