fix poor llama ptb

xvyaward · Mar 7, 2024 · 03cfc99 · 03cfc99
1 parent 27e8523
commit 03cfc99
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/owq/utils/datautils.py b/owq/utils/datautils.py
@@ -3,7 +3,7 @@
 import os
 import torch
 from datasets import load_dataset
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, LlamaTokenizer
 
 def get_wikitext2(nsamples, seed, seqlen, tokenizer, train):
     if train:
@@ -89,6 +89,8 @@ def get_loaders(
     name, nsamples=128, seed=0, seqlen=2048, model='', train=True
 ):
     tokenizer = AutoTokenizer.from_pretrained(model, use_fast=False)
+    if isinstance(tokenizer, LlamaTokenizer) and 'ptb' in name:
+        tokenizer.tokens_trie.data = {}
 
     if 'wikitext2' in name:
         return get_wikitext2(nsamples, seed, seqlen, tokenizer, train)