Repository: EducationalTestingService/rsmtool
Pull request #651 (2023-09-06 21:41:28)

Enhance wandb logging of evaluation metrics
- log all the values of the following tables as run metrics: `consistency`, `eval_short`, `true_score_eval`. This allows easy comparison between runs in the [runs view](https://wandb.ai/etslabs/rsmtool_tutorial/table?workspace=user-tamarlv)

- log confusion matrices as custom charts, this also allows comparing between runs - see in the [workspace view](https://wandb.ai/etslabs/rsmtool_tutorial/workspace?workspace=user-tamarlv)

@desilinguist @mulhod The changes make the tutorial run quite slow for me. Please test and let me know what you think.
to test add the following to `rsmtool` / `rsmeval` tutorial config (change project name accordingly), and run it in this branch:
``` 
    "use_wandb": true,
    "wandb_project": "rsmtool_tutorial",
    "wandb_entity": "etslabs"
```


Patch:
diff --git a/rsmtool/analyzer.py b/rsmtool/analyzer.py
index b8b07e67..5f337dab 100644
--- a/rsmtool/analyzer.py
+++ b/rsmtool/analyzer.py
@@ -20,6 +20,7 @@ from sklearn.metrics import confusion_matrix, mean_squared_error, r2_score
 from skll.metrics import kappa
 
 from .container import DataContainer
+from .utils import wandb
 from .utils.metrics import (
     agreement,
     difference_of_standardized_means,
@@ -1670,7 +1671,7 @@ class Analyzer:
 
         return configuration, DataContainer(datasets=datasets)
 
-    def run_prediction_analyses(self, data_container, configuration):
+    def run_prediction_analyses(self, data_container, configuration, wandb_run=None):
         """
         Run all analyses on the system scores (predictions).
 
@@ -1687,6 +1688,12 @@ class Analyzer:
             following parameters (keys):  {"subgroups", "second_human_score_column",
             "use_scaled_predictions"}.
 
+        wandb_run : wandb.Run
+            The wandb run object if wandb is enabled, ``None`` otherwise.
+            If enabled, all the output data frames will be logged to
+            this run as tables.
+            Defaults to ``None``.
+
         Returns
         -------
         data_container : container.DataContainer
@@ -1790,6 +1797,10 @@ class Analyzer:
         conf_matrix = confusion_matrix(human_scores, system_scores)
         labels = sorted(pd.concat([human_scores, system_scores]).unique())
         df_confmatrix = pd.DataFrame(conf_matrix, index=labels, columns=labels).transpose()
+        # log confusion matrix to W&B
+        wandb.log_confusion_matrix(
+            wandb_run, human_scores, system_scores, "Human-System Confusion Matrix"
+        )
 
         # compute the score distributions of the rounded human and system scores
         df_score_dist = df_preds[["sc1_round", f"{score_type}_trim_round"]].apply(
@@ -1808,6 +1819,9 @@ class Analyzer:
             df_confmatrix_h1h2 = pd.DataFrame(
                 conf_matrix_h1h2, index=labels, columns=labels
             ).transpose()
+            wandb.log_confusion_matrix(
+                wandb_run, human1_scores, human2_scores, "Human1-Human2 Confusion Matrix"
+            )
 
         # Replace any NaNs, which we might get because our model never
         # predicts a particular score label, with zeros.
diff --git a/rsmtool/rsmeval.py b/rsmtool/rsmeval.py
index f83e4d94..80dde80c 100644
--- a/rsmtool/rsmeval.py
+++ b/rsmtool/rsmeval.py
@@ -203,7 +203,9 @@ def run_evaluation(
     (
         pred_analysis_config,
         pred_analysis_data_container,
-    ) = analyzer.run_prediction_analyses(for_pred_data_container, analyzed_config)
+    ) = analyzer.run_prediction_analyses(
+        for_pred_data_container, analyzed_config, wandb_run=wandb_run
+    )
 
     writer.write_experiment_output(
         csvdir, pred_analysis_data_container, reset_index=True, file_format=file_format
diff --git a/rsmtool/rsmtool.py b/rsmtool/rsmtool.py
index 583f0b7a..904273fc 100644
--- a/rsmtool/rsmtool.py
+++ b/rsmtool/rsmtool.py
@@ -333,7 +333,7 @@ def run_experiment(
     (
         pred_analysis_config,
         pred_analysis_data_container,
-    ) = analyzer.run_prediction_analyses(new_pred_data_container, pred_config)
+    ) = analyzer.run_prediction_analyses(new_pred_data_container, pred_config, wandb_run)
 
     # Write out files
     writer.write_experiment_output(
diff --git a/rsmtool/utils/wandb.py b/rsmtool/utils/wandb.py
index 15be1d30..25f11103 100644
--- a/rsmtool/utils/wandb.py
+++ b/rsmtool/utils/wandb.py
@@ -5,7 +5,13 @@ Utility classes and functions for logging to Weights & Biases.
 """
 import wandb
 
-EXCLUDE_WANDB_LOG = ["confMatrix"]
+# excluded dataframes will not be logged as tables or metrics.
+# confusion matrices are logged separately.
+EXCLUDED = ["confMatrix", "confMatrix_h1h2"]
+
+# all values from these dataframes will be logged to the
+# run as metrics in addition to logging as a table artifact.
+METRICS_LOGGED = ["consistency", "eval_short", "true_score_eval"]
 
 
 def init_wandb_run(config_obj):
@@ -54,6 +60,10 @@ def log_dataframe_to_wandb(wandb_run, df, df_name):
     """
     Log a dataframe as a table to W&B if logging to W&B is enabled.
 
+    Dataframes are logged as a table artifact. Values from selected
+    dataframes will also be logged as metrics for easier comparison
+    between runs in the W&B project dashboard.
+
     Parameters
     ----------
     wandb_run : wandb.Run
@@ -63,9 +73,79 @@ def log_dataframe_to_wandb(wandb_run, df, df_name):
     df_name : str
         The name of the dataframe
     """
-    if wandb_run and df_name not in EXCLUDE_WANDB_LOG:
+    if wandb_run and df_name not in EXCLUDED:
         table = wandb.Table(dataframe=df, allow_mixed_types=True)
         wandb_run.log({df_name: table})
+        if df_name in METRICS_LOGGED:
+            indexed_df = df.set_index(df.columns[0])
+            metric_dict = {}
+            for column in indexed_df.columns:
+                col_dict = indexed_df[column].to_dict()
+                metric_dict.update(
+                    {
+                        get_metric_name(df_name, column, row): value
+                        for row, value in col_dict.items()
+                        if value
+                    }
+                )
+            wandb_run.log(metric_dict)
+
+
+def get_metric_name(df_name, col_name, row_name):
+    """
+    Generate the metric name for logging in W&B.
+
+    The name contains the dataframe, column and row names,
+    unless row name is empty or 0 (when dataframe has a single line)
+
+    Parameters
+    ----------
+    df_name : str
+        The dataframe name
+    col_name : str
+        The column name
+    row_name : Union[int,str]
+        The row name, or 0 for some single line dataframes.
+
+    Returns
+    -------
+    metric_name : str
+        The metric name
+    """
+    metric_name = f"{df_name}.{col_name}"
+    if row_name != "" and row_name != 0:
+        metric_name = f"{metric_name}.{row_name}"
+    return metric_name
+
+
+def log_confusion_matrix(wandb_run, human_scores, system_scores, name):
+    """
+    Log a confusion matrix to W&B if logging to W&B is enabled.
+
+    The confusion matrix is added as a custom chart.
+
+    Parameters
+    ----------
+    wandb_run : wandb.Run
+        The wandb run object, or None, if logging to W&B is disabled
+    human_scores : Sequence
+        The human scores for the responses in the data
+    system_scores : Sequence
+        The predicted scores for the responses in the data
+    name : str
+        The chart title
+    """
+    if wandb_run:
+        wandb_run.log(
+            {
+                name: wandb.plot.confusion_matrix(
+                    probs=None,
+                    y_true=human_scores,
+                    preds=system_scores,
+                    title=name,
+                )
+            }
+        )
 
 
 def log_report_to_wandb(wandb_run, report_name, report_path):


Test patch:
diff --git a/tests/test_experiment_rsmtool_4.py b/tests/test_experiment_rsmtool_4.py
index c40ad233..4d3db2c4 100644
--- a/tests/test_experiment_rsmtool_4.py
+++ b/tests/test_experiment_rsmtool_4.py
@@ -164,13 +164,15 @@ class TestExperimentRsmtool4(unittest.TestCase):
         with self.assertRaises(ValueError):
             run_experiment(config, output_dir, overwrite_output=True)
 
-    def test_run_experiment_with_wandb(self):
+    @patch("wandb.init")
+    @patch("wandb.plot.confusion_matrix")
+    def test_run_experiment_with_wandb(self, mock_plot_conf_mat, mock_wandb_init):
         source = "wandb"
         experiment_id = "wandb"
         config_file = join(rsmtool_test_dir, "data", "experiments", source, f"{experiment_id}.json")
         mock_wandb_run = Mock()
-        with patch("wandb.init") as mock_wandb_init:
-            mock_wandb_init.return_value = mock_wandb_run
-            do_run_experiment(source, experiment_id, config_file)
+        mock_wandb_init.return_value = mock_wandb_run
+        do_run_experiment(source, experiment_id, config_file)
         mock_wandb_init.assert_called_with(project="wandb_project", entity="wandb_entity")
         mock_wandb_run.log_artifact.assert_called_once()
+        mock_plot_conf_mat.assert_called()
diff --git a/tests/test_wandb.py b/tests/test_wandb.py
index 8cd26c6b..2cf6196a 100644
--- a/tests/test_wandb.py
+++ b/tests/test_wandb.py
@@ -7,6 +7,7 @@ import pandas as pd
 
 from rsmtool.configuration_parser import Configuration
 from rsmtool.utils.wandb import (
+    get_metric_name,
     init_wandb_run,
     log_configuration_to_wandb,
     log_dataframe_to_wandb,
@@ -118,3 +119,12 @@ class TestWandb(unittest.TestCase):
             # assert that no calls to wandb objects have been made
             mock_wandb_artifact.assert_not_called()
             mock_report_artifact.add_file.assert_not_called()
+
+    def test_get_metric_name(self):
+        """Test the method get_metric_name."""
+        # with a valid row name
+        self.assertEquals("df.col.row", get_metric_name("df", "col", "row"))
+        # with an empty row name
+        self.assertEquals("df.col", get_metric_name("df", "col", ""))
+        # with 0 as row name
+        self.assertEquals("df.col", get_metric_name("df", "col", 0))