Update example

neuralmagic · Jul 18, 2024 · be6eef2 · be6eef2
1 parent ab3dad3
commit be6eef2
Showing 1 changed file with 1 addition and 8 deletions.
diff --git a/example_dataset.py b/example_dataset.py
@@ -9,17 +9,10 @@
 tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=True)
 tokenizer.pad_token = tokenizer.eos_token
 
-MAX_SEQUENCE_LENGTH = 2048
 ds = load_dataset("mgoin/ultrachat_2k", split="train_sft").select(range(512))
 def preprocess(example):
     example = tokenizer.apply_chat_template(example["messages"], tokenize=False)
-    return tokenizer(
-        example,
-        padding=False,
-        max_length=MAX_SEQUENCE_LENGTH,
-        truncation=True,
-        add_special_tokens=False,
-    )
+    return tokenizer(example, max_length=2048, truncation=True, add_special_tokens=False)
 ds = ds.map(preprocess, remove_columns=ds.column_names)
 
 quantize_config = BaseQuantizeConfig(quant_method="fp8", activation_scheme="static")