Update engine.py

engine.py

@@ -9,6 +9,7 @@ from functools import partial
 from typing import Generator, Optional, List, Dict, Any, Tuple
 from datasets import Dataset, load_dataset
 from trl import SFTConfig, SFTTrainer
+from trl.trainer.utils import remove_none_values
 from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl
 from huggingface_hub import HfApi, model_info, metadata_update
 
@@ -374,6 +375,10 @@ class FunctionGemmaEngine:
         return fig
 
     def _evaluate_model(self, test_dataset) -> Generator[str, None, None]:
+        # Tabular backends like Arrow/Parquet insert `None` for mismatched keys in nested structures. Clean them from
+        # sampled data.
+        test_dataset = test_dataset.with_transform(remove_none_values) 
+
         results = []
         success_count = 0
         for idx, item in enumerate(test_dataset):