From ef651247d0daaafcab0d8f1c1ce7e3838bcd95a2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?DAN=E2=84=A2?= <dranger003@gmail.com>
Date: Sat, 2 Mar 2024 21:56:45 -0500
Subject: [PATCH] Support special tokens as reverse/anti prompt.

---
 examples/main/main.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/examples/main/main.cpp b/examples/main/main.cpp
index 34e84d0d4..861a88d58 100644
--- a/examples/main/main.cpp
+++ b/examples/main/main.cpp
@@ -760,7 +760,8 @@ int main(int argc, char ** argv) {
                         ? last_output.length() - static_cast<size_t>(antiprompt.length() + extra_padding)
                         : 0;
 
-                    if (last_output.find(antiprompt, search_start_pos) != std::string::npos) {
+                    auto tmp = ::llama_tokenize(ctx, antiprompt, false, true);
+                    if (last_output.find(antiprompt, search_start_pos) != std::string::npos || (tmp.size() == 1 && llama_sampling_last(ctx_sampling) == tmp[0])) {
                         if (params.interactive) {
                             is_interacting = true;
                         }