|
38 | 38 | #include "marian_conv.cpp" |
39 | 39 | #undef main |
40 | 40 |
|
| 41 | +#include <map> |
| 42 | +#include <string> |
| 43 | +#include <tuple> |
41 | 44 | #include "3rd_party/ExceptionWithCallStack.h" |
| 45 | +#include "3rd_party/spdlog/details/format.h" |
42 | 46 |
|
43 | 47 | int main(int argc, char** argv) { |
44 | 48 | using namespace marian; |
| 49 | + using MainFunc = std::function<int(int, char**)>; |
45 | 50 |
|
46 | | - if(argc > 1 && argv[1][0] != '-') { |
| 51 | + std::map<std::string, std::tuple<MainFunc, std::string>> subcmds = { |
| 52 | + {"train", {mainTrainer, "Train a model (default)"}}, |
| 53 | + {"decode", {mainDecoder, "Decode or translate text"}}, |
| 54 | + {"score", {mainScorer, "Score translations"}}, |
| 55 | + {"embed", {mainEmbedder, "Embed text"}}, |
| 56 | + {"evaluate", {mainEvaluator, "Run Evaluator metric"}}, |
| 57 | + {"vocab", {mainVocab, "Create vocabulary"}}, |
| 58 | + {"convert", {mainConv, "Convert model file format"}} |
| 59 | + }; |
| 60 | + // no arguments, or the first arg is "?"", print help message |
| 61 | + if (argc == 1 || (argc == 2 && (std::string(argv[1]) == "?") )) { |
| 62 | + std::cout << "Usage: " << argv[0] << " COMMAND [ARGS]" << std::endl; |
| 63 | + std::cout << "Commands:" << std::endl; |
| 64 | + for (auto&& [name, val] : subcmds) { |
| 65 | + std::cerr << fmt::format("{:10} : {}\n", name, std::get<1>(val)); |
| 66 | + } |
| 67 | + return 0; |
| 68 | + } |
| 69 | + |
| 70 | + if (argc > 1 && argv[1][0] != '-') { |
47 | 71 | std::string cmd = argv[1]; |
48 | 72 | argc--; |
49 | 73 | argv[1] = argv[0]; |
50 | 74 | argv++; |
51 | | - if(cmd == "train") return mainTrainer(argc, argv); |
52 | | - else if(cmd == "decode") return mainDecoder(argc, argv); |
53 | | - else if (cmd == "score") return mainScorer(argc, argv); |
54 | | - else if (cmd == "embed") return mainEmbedder(argc, argv); |
55 | | - else if (cmd == "evaluate") return mainEvaluator(argc, argv); |
56 | | - else if (cmd == "vocab") return mainVocab(argc, argv); |
57 | | - else if (cmd == "convert") return mainConv(argc, argv); |
58 | | - std::cerr << "Command must be train, decode, score, embed, vocab, or convert." << std::endl; |
59 | | - exit(1); |
60 | | - } else |
| 75 | + // check if the command is known, and if so, call the respective function |
| 76 | + // c++20 has contains() for maps, but we're not there yet, so we use count() |
| 77 | + if (subcmds.count(cmd) > 0) { |
| 78 | + auto [func, desc] = subcmds[cmd]; |
| 79 | + return func(argc, argv); |
| 80 | + } |
| 81 | + else { |
| 82 | + std::cerr << "Unknown command: " << cmd << ". Known commands are:" << std::endl; |
| 83 | + for (auto&& [name, val] : subcmds) { |
| 84 | + std::cerr << fmt::format("{:10} : {}\n", name, std::get<1>(val)); |
| 85 | + } |
| 86 | + return 1; |
| 87 | + } |
| 88 | + } |
| 89 | + else |
61 | 90 | return mainTrainer(argc, argv); |
62 | 91 | } |
0 commit comments