diff options
Diffstat (limited to 'lm/query_main.cc')
-rw-r--r-- | lm/query_main.cc | 51 |
1 files changed, 37 insertions, 14 deletions
diff --git a/lm/query_main.cc b/lm/query_main.cc index 27d3a1a56..bd4fde62f 100644 --- a/lm/query_main.cc +++ b/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); |