Update README.md

README.md

@@ -116,11 +116,13 @@ Although AutoGPTQ was used for this particular model, Neural Magic is transition
 ```python
 from transformers import AutoTokenizer
 from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
+from datasets import load_dataset
 import random
 
 model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
 
 num_samples = 512
+max_seq_len = 4096
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
@@ -131,7 +133,7 @@ dataset = load_dataset(dataset_name, split="train")
 ds = dataset.shuffle().select(range(num_samples))
 ds = ds.map(preprocess_fn)
 
-return [
+examples = [
     tokenizer(
         example["text"], padding=False, max_length=max_seq_len, truncation=True,
     ) for example in ds