New calibration

calibrate.py

@@ -42,34 +42,17 @@ tokenizer = AutoTokenizer.from_pretrained(model_stub)
 
 output_dir = f"/{model_name}-FP8"
 
-DATASET_ID = "HuggingFaceH4/ultrachat_200k"
-DATASET_SPLIT = "train_sft"
+DATASET_ID = "neuralmagic/LLM_compression_calibration"
+DATASET_SPLIT = "train"
 NUM_CALIBRATION_SAMPLES = 512
-MAX_SEQUENCE_LENGTH = 4096
+MAX_SEQUENCE_LENGTH = 8192
 
-ds = load_dataset(DATASET_ID, split=DATASET_SPLIT)
-ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
-
-def preprocess(example):
-    return {
-        "text": tokenizer.apply_chat_template(
-            example["messages"],
-            tokenize=False,
-        )
-    }
-
-ds = ds.map(preprocess)
+def preprocess_fn(example):
+  return {"text": tokenizer.apply_chat_template(example["messages"], add_generation_prompt=False, tokenize=False)}
 
-def tokenize(sample):
-    return tokenizer(
-        sample["text"],
-        padding=False,
-        max_length=MAX_SEQUENCE_LENGTH,
-        truncation=True,
-        add_special_tokens=False,
-    )
-
-ds = ds.map(tokenize, remove_columns=ds.column_names)
+ds = load_dataset(DATASET_ID, split=DATASET_SPLIT)
+ds = ds.shuffle().select(range(NUM_CALIBRATION_SAMPLES))
+ds = ds.map(preprocess_fn)
 
 oneshot(
     model=model,

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fcbde45e4f901efddc4005b5aa56dc04bc648e24a2b40763484ae5e97da77d28
+oid sha256:e164b5b6ae0c8a2e3fc8f8684513e31a79be4285d917cd1e3b9787272e8d8a83
 size 903038064