Update example_mixtral.py

neuralmagic · Jun 27, 2024 · 1958d07 · 1958d07
1 parent 2a9330c
commit 1958d07
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/example_mixtral.py b/examples/example_mixtral.py
@@ -9,7 +9,7 @@
 tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=True)
 tokenizer.pad_token = tokenizer.eos_token
 
-ds = load_dataset("mgoin/ultrachat_2k", split="train_sft").select(range(10))
+ds = load_dataset("mgoin/ultrachat_2k", split="train_sft")
 examples = [tokenizer.apply_chat_template(batch["messages"], tokenize=False) for batch in ds]
 examples = tokenizer(examples, padding=True, truncation=True, return_tensors="pt").to("cuda")