clean-up

Author: dsikka

tests/e2e/e2e_utils.py

@@ -1,14 +1,15 @@
+from typing import Callable
+
 import torch
 import transformers
 from datasets import load_dataset
 from loguru import logger
-from transformers import AutoProcessor
+from transformers import AutoProcessor, DefaultDataCollator
 
 from llmcompressor import oneshot
 from llmcompressor.modifiers.quantization import GPTQModifier, QuantizationModifier
 from tests.test_timer.timer_utils import log_time
 from tests.testing_utils import process_dataset
-from transformers import DefaultDataCollator
 
 
 def load_model(model: str, model_class: str, device_map: str | None = None):
@@ -35,9 +36,12 @@ def run_oneshot_for_e2e_testing(
     dataset_config: str,
     scheme: str,
     quant_type: str,
+    shuffle_calibration_samples: bool = True,
+    data_collator: str | Callable = DefaultDataCollator(),
 ):
     # Load model.
     oneshot_kwargs = {}
+    oneshot_kwargs["data_collator"] = data_collator
 
     loaded_model = load_model(model=model, model_class=model_class)
     processor = AutoProcessor.from_pretrained(model)
@@ -75,6 +79,7 @@ def data_collator(batch):
             oneshot_kwargs["data_collator"] = data_collator
 
     oneshot_kwargs["model"] = loaded_model
+    oneshot_kwargs["shuffle_calibration_samples"] = shuffle_calibration_samples
     if recipe:
         oneshot_kwargs["recipe"] = recipe
     else:
@@ -95,11 +100,8 @@ def data_collator(batch):
             )
 
     # Apply quantization.
-
+    breakpoint()
     logger.info("ONESHOT KWARGS", oneshot_kwargs)
-
-    oneshot_kwargs["shuffle_calibration_samples"] = True
-    oneshot_kwargs["data_collator"] = DefaultDataCollator()
     _run_oneshot(**oneshot_kwargs)
 
     return oneshot_kwargs["model"], processor