diff --git a/examples/lookahead/README.md b/examples/lookahead/README.md index a69a471b47d39..61f9840baccc0 100644 --- a/examples/lookahead/README.md +++ b/examples/lookahead/README.md @@ -5,3 +5,5 @@ Demonstration of lookahead decoding technique: https://lmsys.org/blog/2023-11-21-lookahead-decoding/ More info: https://github.com/ggerganov/llama.cpp/pull/4207 + +To set W, N and G manually, set them as environment variables. diff --git a/examples/lookahead/lookahead.cpp b/examples/lookahead/lookahead.cpp index e2551e7a494c2..ac0812258516a 100644 --- a/examples/lookahead/lookahead.cpp +++ b/examples/lookahead/lookahead.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include @@ -34,6 +35,17 @@ struct ngram_container { std::vector tokens; }; +int int_from_env(const char *name, int default_value) { + char *value = std::getenv(name); + if (value != nullptr) { + int v = atoi(value); + v = v ? v : default_value; + return v; + } else { + return default_value; + } +} + int main(int argc, char ** argv) { gpt_params params; @@ -41,9 +53,10 @@ int main(int argc, char ** argv) { return 1; } - const int W = 15; // lookahead window - const int N = 5; // n-gram size - const int G = 15; // max verification n-grams + int W = int_from_env("W",15); // lookahead window + int N = int_from_env("N",5); // n-gram size + int G = int_from_env("G",15); // max verification n-grams + const bool dump_kv_cache = params.dump_kv_cache;