1 files changed, 18 insertions, 15 deletions

diff --git a/klm/lm/ngram_query.cc b/klm/lm/ngram_query.cc
index d6da02e3..9454a6d1 100644
--- a/klm/lm/ngram_query.cc
+++ b/klm/lm/ngram_query.cc
@@ -35,14 +35,14 @@ void PrintUsage(const char *message) {
   }
 }
 
-template <class Model> void Query(const Model &model) {
+template <class Model> void Query(const Model &model, bool sentence_context) {
   PrintUsage("Loading statistics:\n");
   typename Model::State state, out;
   lm::FullScoreReturn ret;
   std::string word;
 
   while (std::cin) {
-    state = model.BeginSentenceState();
+    state = sentence_context ? model.BeginSentenceState() : model.NullContextState();
     float total = 0.0;
     bool got = false;
     unsigned int oov = 0;
@@ -52,7 +52,7 @@ template <class Model> void Query(const Model &model) {
       if (vocab == 0) ++oov;
       ret = model.FullScore(state, vocab, out);
       total += ret.prob;
-      std::cout << word << '=' << vocab << ' ' << static_cast<unsigned int>(ret.ngram_length)  << ' ' << ret.prob << '\n';
+      std::cout << word << '=' << vocab << ' ' << static_cast<unsigned int>(ret.ngram_length)  << ' ' << ret.prob << '\t';
       state = out;
       char c;
       while (true) {
@@ -67,9 +67,11 @@ template <class Model> void Query(const Model &model) {
       if (c == '\n') break;
     }
     if (!got && !std::cin) break;
-    ret = model.FullScore(state, model.GetVocabulary().EndSentence(), out);
-    total += ret.prob;
-    std::cout << "</s>=" << model.GetVocabulary().EndSentence() << ' ' << static_cast<unsigned int>(ret.ngram_length)  << ' ' << ret.prob << '\n';
+    if (sentence_context) {
+      ret = model.FullScore(state, model.GetVocabulary().EndSentence(), out);
+      total += ret.prob;
+      std::cout << "</s>=" << model.GetVocabulary().EndSentence() << ' ' << static_cast<unsigned int>(ret.ngram_length)  << ' ' << ret.prob << '\t';
+    }
     std::cout << "Total: " << total << " OOV: " << oov << '\n';
   }
   PrintUsage("After queries:\n");
@@ -82,29 +84,30 @@ template <class Model> void Query(const char *name) {
 }
 
 int main(int argc, char *argv[]) {
-  if (argc < 2) {
-    std::cerr << "Pass language model name." << std::endl;
-    return 0;
+  if (!(argc == 2 || (argc == 3 && !strcmp(argv[2], "null")))) {
+    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 = (argc == 2);
   lm::ngram::ModelType model_type;
   if (lm::ngram::RecognizeBinary(argv[1], model_type)) {
     switch(model_type) {
       case lm::ngram::HASH_PROBING:
-        Query<lm::ngram::ProbingModel>(argv[1]);
-        break;
-      case lm::ngram::HASH_SORTED:
-        Query<lm::ngram::SortedModel>(argv[1]);
+        Query<lm::ngram::ProbingModel>(argv[1], sentence_context);
         break;
       case lm::ngram::TRIE_SORTED:
-        Query<lm::ngram::TrieModel>(argv[1]);
+        Query<lm::ngram::TrieModel>(argv[1], sentence_context);
         break;
+      case lm::ngram::HASH_SORTED:
       default:
         std::cerr << "Unrecognized kenlm model type " << model_type << std::endl;
         abort();
     }
   } else {
-    Query<lm::ngram::ProbingModel>(argv[1]);
+    Query<lm::ngram::ProbingModel>(argv[1], sentence_context);
   }
 
   PrintUsage("Total time including destruction:\n");
+  return 0;
 }