Skip to content

Commit edcc432

Browse files
committed
Add data path to runner
^ Differential Revision: [D73617661](https://our.internmc.facebook.com/intern/diff/D73617661/) [ghstack-poisoned]
1 parent 5f59b76 commit edcc432

File tree

1 file changed

+12
-1
lines changed

1 file changed

+12
-1
lines changed

examples/models/llama/main.cpp

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,12 @@ DEFINE_string(
2020
"llama2.pte",
2121
"Model serialized in flatbuffer format.");
2222

23+
DEFINE_string(
24+
data_path,
25+
"",
26+
"Data file for the model."
27+
);
28+
2329
DEFINE_string(tokenizer_path, "tokenizer.bin", "Tokenizer stuff.");
2430

2531
DEFINE_string(prompt, "The answer to the ultimate question is", "Prompt.");
@@ -49,6 +55,11 @@ int32_t main(int32_t argc, char** argv) {
4955
// and users can create their own DataLoaders to load from arbitrary sources.
5056
const char* model_path = FLAGS_model_path.c_str();
5157

58+
std::optional<std::string> data_path = std::nullopt;
59+
if (!FLAGS_data_path.empty()) {
60+
data_path = FLAGS_data_path.c_str();
61+
}
62+
5263
const char* tokenizer_path = FLAGS_tokenizer_path.c_str();
5364

5465
const char* prompt = FLAGS_prompt.c_str();
@@ -74,7 +85,7 @@ int32_t main(int32_t argc, char** argv) {
7485
#endif
7586
// create llama runner
7687
// @lint-ignore CLANGTIDY facebook-hte-Deprecated
77-
example::Runner runner(model_path, tokenizer_path);
88+
example::Runner runner(model_path, tokenizer_path, data_path);
7889

7990
if (warmup) {
8091
// @lint-ignore CLANGTIDY facebook-hte-Deprecated

0 commit comments

Comments
 (0)