Add metrics and aggregators for ARC AGI pipeline

Author: ilmarinen

eureka_ml_insights/data_utils/arc_agi_utils.py

@@ -0,0 +1,39 @@
+import re
+from dataclasses import dataclass
+
+import pandas as pd
+
+from .transform import DFTransformBase
+
+
+@dataclass
+class ARCAGI_ExtractAnswer(DFTransformBase):
+    model_output_column: str
+    model_answer_column: str
+
+    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
+        df[self.model_answer_column] = df[self.model_output_column].apply(self.parse_output_answer)
+        return df
+
+    @staticmethod
+    def parse_output_answer(response):
+        """
+        Parse the input string to extract answer of a given ARCAGI question.
+        Parameters:
+            response (str): Input string containing answer X in the form of "<output>final answer string</output>".
+        Returns: 
+            answer (str): The final answer string with leading and training spaces stripped.
+        """
+        answer = ""
+
+        if response is None:
+            return ""
+        elif response.find("<output>") == -1 or response.find("</output>") == -1:
+            return ""
+
+        start_index = response.find("<output>") + len("<output>")
+        end_index = response.find("</output>")
+
+        answer = response[start_index:end_index].strip()
+
+        return answer

eureka_ml_insights/user_configs/arc_agi.py

@@ -4,16 +4,17 @@
 from eureka_ml_insights.core import Inference, PromptProcessing
 from eureka_ml_insights.core.data_processing import DataProcessing
 from eureka_ml_insights.core.eval_reporting import EvalReporting
-from eureka_ml_insights.data_utils.ba_calendar_utils import (
-    BA_Calendar_ExtractAnswer,
+from eureka_ml_insights.data_utils.arc_agi_utils import (
+    ARCAGI_ExtractAnswer,
 )
 from eureka_ml_insights.data_utils.data import (
     DataLoader,
     DataReader,
     HFDataReader,
 )
-from eureka_ml_insights.metrics.ba_calendar_metrics import BACalendarMetric
+from eureka_ml_insights.metrics.metrics_base import ExactMatch
 from eureka_ml_insights.metrics.reports import (
+    CountAggregator,
     AverageAggregator,
     BiLevelCountAggregator,
     BiLevelAggregator,
@@ -84,11 +85,62 @@ def configure_pipeline(self, model_config=None, resume_from=None, resume_logdir=
         if resume_logdir:
             self.log_dir = resume_from.split("/")[0:len(resume_from.split("/")) - 1]
 
+        # Configure the evaluation and reporting component for evaluation and dataset level aggregation
+        self.evalreporting_comp = EvalReportingConfig(
+            component_type=EvalReporting,
+            data_reader_config=DataSetConfig(
+                DataReader,
+                {
+                    "path": os.path.join(self.inference_comp.output_dir, "inference_result.jsonl"),
+                    "format": ".jsonl",
+                    "transform": SequenceTransform(
+                        [
+                            ExtractUsageTransform(model_config),
+                            ColumnRename(
+                                name_mapping={
+                                    "model_output": "raw_output",
+                                }
+                            ),
+                            AddColumn("model_output"),
+                            ARCAGI_ExtractAnswer("raw_output", "model_output"),
+                        ]
+                    ),
+                },
+            ),
+            metric_config=MetricConfig(ExactMatch),
+            aggregator_configs=[
+                AggregatorConfig(
+                    CountAggregator,
+                    {
+                        "column_names": [
+                            "ExactMatch_result",
+                        ],
+                        "filename_base": "OverallMetrics_Separate_Runs_Grouped",
+                        "normalize": True,
+                        "group_by": "split",
+                    },
+                ),
+                # the next three reports take the average and std for all repeats
+                # the resulting numbers are the average and std of N pass@1 scores, where N is number of repeats
+                AggregatorConfig(
+                    CountAggregator, 
+                    {
+                        "column_names": [
+                            "ExactMatch_result",
+                        ],
+                        "normalize": True,
+                        "filename_base": "OverallMetrics_Separate_Runs_Total",
+                    }),
+            ],
+            output_dir=os.path.join(self.log_dir, "eval_report"),
+        )
+
         # Configure the pipeline
         return PipelineConfig(
             [
                 self.data_processing_comp,
                 self.inference_comp,
+                self.evalreporting_comp,
             ],
             self.log_dir,
         )