From d7747aa32e27f7f708309fa4fc675ee78fe86835 Mon Sep 17 00:00:00 2001
From: Ivan Chikish <nekotekina@gmail.com>
Date: Mon, 26 Feb 2024 03:26:51 +0300
Subject: [PATCH] examples/main: minor token print filtering for LLaVA

In ChatML mode, for Yi-34B llava-1.6, don't print <|im_end|>
---
 examples/main/main.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/examples/main/main.cpp b/examples/main/main.cpp
index 34e84d0d42f878..c3124e8ff134de 100644
--- a/examples/main/main.cpp
+++ b/examples/main/main.cpp
@@ -725,7 +725,9 @@ int main(int argc, char ** argv) {
         // display text
         if (input_echo && display) {
             for (auto id : embd) {
-                const std::string token_str = llama_token_to_piece(ctx, id);
+                std::string token_str = llama_token_to_piece(ctx, id);
+                if (params.chatml && token_str == "<|im_end|>")
+                    token_str.clear();
                 printf("%s", token_str.c_str());
 
                 if (embd.size() > 1) {