Просмотр исходного кода

Fix LlamaTokenizer eos_token (attempt)

oobabooga 2 лет назад
Родитель
Сommit
a3085dba07
1 измененных файлов с 3 добавлено и 0 удалено
  1. 3 0
      modules/models.py

+ 3 - 0
modules/models.py

@@ -174,6 +174,9 @@ def load_model(model_name):
         tokenizer = AutoTokenizer.from_pretrained(Path(f"{shared.args.model_dir}/gpt-j-6B/"))
     elif type(model) is transformers.LlamaForCausalLM:
         tokenizer = LlamaTokenizer.from_pretrained(Path(f"{shared.args.model_dir}/{shared.model_name}/"), clean_up_tokenization_spaces=True)
+        tokenizer.eos_token_id = 2
+        tokenizer.bos_token_id = 1
+        tokenizer.pad_token_id = 0
     else:
         tokenizer = AutoTokenizer.from_pretrained(Path(f"{shared.args.model_dir}/{shared.model_name}/"))
     tokenizer.truncation_side = 'left'