aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--examples/main/main.cpp6
1 files changed, 6 insertions, 0 deletions
diff --git a/examples/main/main.cpp b/examples/main/main.cpp
index c7c5915..6131f5b 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);
}
}