Keep GALACTICA special tokens (oobabooga#300)

TheTerrasque · Mar 16, 2023 · a577fb1 · a577fb1
1 parent 25a00ea
commit a577fb1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/modules/text_generation.py b/modules/text_generation.py
@@ -38,7 +38,7 @@ def encode(prompt, tokens_to_generate=0, add_special_tokens=True):
 
 def decode(output_ids):
     # Open Assistant relies on special tokens like <|endoftext|>
-    if re.match('oasst-*', shared.model_name.lower()):
+    if re.match('(oasst|galactica)-*', shared.model_name.lower()):
         return shared.tokenizer.decode(output_ids, skip_special_tokens=False)
     else:
         reply = shared.tokenizer.decode(output_ids, skip_special_tokens=True)