From 248367605ead6fb7c36d2bfb1ebd8f00a23f7c71 Mon Sep 17 00:00:00 2001 From: DannyDaemonic Date: Mon, 29 May 2023 05:13:40 -0700 Subject: [PATCH] Work around for recalculating logits in cached prompts (Fixes #1585) (#1609) * Work around for recalculating logits in cached prompts --- examples/main/main.cpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/examples/main/main.cpp b/examples/main/main.cpp index c7c591537..6131f5b46 100644 --- a/examples/main/main.cpp +++ b/examples/main/main.cpp @@ -360,6 +360,12 @@ int main(int argc, char ** argv) { } } if (i > 0) { + // check if we've used up all the prompt but not all cached tokens + if (embd.size() == i && n_session_consumed < (int) session_tokens.size()) { + // force revaluation of the last token to recalculate logits + i--; + n_past--; + } embd.erase(embd.begin(), embd.begin() + i); } }