From 643b7df0c1b3c19f5ba523b8aa3ad478e7bca255 Mon Sep 17 00:00:00 2001
From: "upskyy (Patrick Ha)" <patrick@rtzr.ai>
Date: Thu, 11 Jul 2024 19:13:58 +0900
Subject: [PATCH] Add gemma2 stop tokens

---
 generator.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/generator.py b/generator.py
index ef9753e..1302fad 100644
--- a/generator.py
+++ b/generator.py
@@ -44,7 +44,7 @@
     temperature=0,
     skip_special_tokens=True,
     max_tokens=args.model_len,
-    stop=["<|endoftext|>", "[INST]", "[/INST]", "<|im_end|>", "<|end|>", "<|eot_id|>"],
+    stop=["<|endoftext|>", "[INST]", "[/INST]", "<|im_end|>", "<|end|>", "<|eot_id|>", "<end_of_turn>", "<eos>"],
 )
 
 df_questions = pd.read_json("questions.jsonl", orient="records", encoding="utf-8-sig", lines=True)