works for vision

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

examples/multimodal_vision/gemma3_example.py

@@ -1,7 +1,7 @@
 import requests
-import torch
 from PIL import Image
 from transformers import AutoProcessor, Gemma3ForConditionalGeneration
+from transformers.data import DataCollatorWithPadding
 
 from llmcompressor import oneshot
 from llmcompressor.modifiers.quantization import GPTQModifier
@@ -11,6 +11,7 @@
 model_id = "google/gemma-3-4b-it"
 model = Gemma3ForConditionalGeneration.from_pretrained(model_id, torch_dtype="auto")
 processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+collator = DataCollatorWithPadding(processor.tokenizer)
 
 # Oneshot arguments
 DATASET_ID = "flickr30k"
@@ -19,10 +20,10 @@
 MAX_SEQUENCE_LENGTH = 2048
 
 
-# Define a oneshot data collator for multimodal inputs.
-def data_collator(batch):
-    assert len(batch) == 1
-    return {key: torch.tensor(value) for key, value in batch[0].items()}
+def data_collator(features: list[dict[str, object]]):
+    # remove extra dim added by vision processor
+    features = [{key: feature[key][0] for key in feature} for feature in features]
+    return collator(features)
 
 
 # Recipe
@@ -48,7 +49,8 @@ def data_collator(batch):
     max_seq_length=MAX_SEQUENCE_LENGTH,
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
     trust_remote_code_model=True,
-    # data_collator=data_collator,
+    batch_size=4,
+    data_collator=data_collator,
 )
 
 # Confirm generations of the quantized model look sane.