Add dataset metadata handling in training pipeline

train_model.py

@@ -22,10 +22,12 @@ def run_grievance_training_pipeline():
 
 
         print(f"[{time.strftime('%H:%M:%S')}] Loading dataset from hub: {configs.dataset_repo_id} ...", flush=True)
-        dataset = load_dataset_from_hub(
+        data = load_dataset_from_hub(
             model_repo=configs.dataset_repo_id,
             hf_token=configs.hf_token
         )
+        dataset = data['dataset']
+        dataset_metadata = data['metadata']
 
         # Print dataset splits and sizes if available
         def _safe_len(split):
@@ -55,6 +57,7 @@ def run_grievance_training_pipeline():
             train_dataset=dataset['train'],
             eval_dataset=dataset['eval'],
             test_dataset=dataset['test'],
+            dataset_metadata= dataset_metadata, 
             space_repo_id=configs.space_repo_id,
             hf_training_args={"hub_model_id": configs.hub_model_id},
             api_endpoint=configs.api_endpoint,