examples/main: minor token print filtering for LLaVA

In ChatML mode, for Yi-34B llava-1.6, don't print <|im_end|>
ggerganov · Feb 27, 2024 · d7747aa · d7747aa
1 parent ce85323
commit d7747aa
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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) {