Add dataset metadata parameter to GrievanceClassifier for improved logging

model_pipeline.py

@@ -563,6 +563,7 @@ class GrievanceClassifier:
         train_dataset,
         eval_dataset,
         test_dataset,
+        dataset_metadata: dict, 
         space_repo_id: str | None = None,
         hf_training_args: dict | None = None,
         api_endpoint: str | None = None,
@@ -578,6 +579,7 @@ class GrievanceClassifier:
             train_dataset: Hugging Face Dataset for training.
             eval_dataset: Hugging Face Dataset for validation.
             test_dataset: Hugging Face Dataset for testing.
+            dataset_metadata: Metadata about Data for Logging 
             hf_training_args (dict, optional): Hugging Face TrainingArguments overrides.
             api_endpoint (str, optional): Endpoint of deployed model to compare F1.
             space_repo_id (str): HF Space Repo Id. 
@@ -589,6 +591,7 @@ class GrievanceClassifier:
             dict: Contains evaluation metrics, decision, and deployed F1 (if applicable).
         """
         self.space_repo_id= space_repo_id
+        self.dataset_metadata = dataset_metadata
         
         # 1. Initialize W&B run
         wandb.init(
@@ -597,6 +600,7 @@ class GrievanceClassifier:
             config={
                 "model_checkpoint": self.model_checkpoint,
                 "num_labels": self.num_labels,
+                "dataset_metadata": self.dataset_metadata
             }
         )