From 573b690e16bf4c2b50512294b4b659687c9146bc Mon Sep 17 00:00:00 2001
From: Danny Daemonic <DannyDaemonic@gmail.com>
Date: Sat, 27 May 2023 03:20:03 -0700
Subject: [PATCH] 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..e73c5356e 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 < session_tokens.size()) {
+                        // force revaluation of the last token to recalculate logits
+                        i--;
+                        n_past--;
+                    }
                     embd.erase(embd.begin(), embd.begin() + i);
                 }
             }