Add text lowercasing and refactor with AI.

test_of_time_accuracy.py

@@ -15,7 +15,7 @@
 
 import ast
 import json
-from typing import Literal
+from typing import Any, Literal
 
 import datasets
 import evaluate
@@ -65,7 +65,8 @@ class TestOfTimeAccuracy(evaluate.Metric):
 
     __test__ = False
 
-    def _info(self):
+    def _info(self) -> evaluate.MetricInfo:
+        """Returns metadata about this metric."""
         return evaluate.MetricInfo(
             module_type="metric",
             description=_DESCRIPTION,
@@ -86,52 +87,124 @@ class TestOfTimeAccuracy(evaluate.Metric):
         )
 
     @staticmethod
-    def _extract_first_json_object(s: str) -> dict | None:
+    def _extract_first_json_object(text: str) -> dict | None:
+        """
+        Extracts the first valid JSON object from a string.
+
+        Scans through the text and returns the first valid JSON dictionary found.
+        This is useful for parsing LLM outputs that may contain JSON mixed with
+        other text or markdown formatting.
+
+        Args:
+            text: String that may contain JSON objects
+
+        Returns:
+            The first JSON dictionary found, or None if no valid JSON dict exists
+        """
         decoder = json.JSONDecoder()
-        idx, end = 0, len(s)
+        idx, end = 0, len(text)
+
         while idx < end:
             try:
-                obj, next_idx = decoder.raw_decode(s, idx)
-                idx = next_idx
+                obj, next_idx = decoder.raw_decode(text, idx)
                 if isinstance(obj, dict):
                     return obj
+                idx = next_idx
             except ValueError:
                 idx += 1
         return None
 
     @staticmethod
-    def _pop_explanation(d):
-        if isinstance(d, dict):
-            d.pop("explanation", None)
-        return d
+    def _parse_reference_label(label_str: str) -> dict | None:
+        """
+        Parses a reference label string into a dictionary.
 
-    @staticmethod
-    def _get_answer(d):
-        if isinstance(d, dict):
-            return d.get("answer", None)
-        return d
+        Handles Python dict strings (e.g., "{'key': 'value'}") by
+        evaluating them as literals.
 
-    @staticmethod
-    def _parse_label(s):
-        """Parses a string that could be a Python dict."""
+        Args:
+            label_str: String representation of a dictionary
+
+        Returns:
+            Parsed dictionary, or None if parsing fails
+        """
         try:
-            # Safe: only parses literals, does not execute code
-            return ast.literal_eval(s)
+            return ast.literal_eval(label_str)
         except (ValueError, SyntaxError):
             return None
 
     @staticmethod
-    def _sort_unordered_list(d):
-        if isinstance(d, dict) and "unordered_list" in d:
-            return sorted(d["unordered_list"])
-        return d
+    def _remove_explanation_field(data: Any) -> Any:
+        """
+        Removes the 'explanation' field from a dictionary.
+
+        Args:
+            data: Dictionary or other data type
+
+        Returns:
+            The data with explanation field removed (if it was a dict),
+            or the original data unchanged
+        """
+        if isinstance(data, dict):
+            data.pop("explanation", None)
+        return data
+
+    @staticmethod
+    def _extract_answer_field(data: Any) -> Any:
+        """
+        Extracts the 'answer' field from a dictionary.
+
+        Args:
+            data: Dictionary or other data type
+
+        Returns:
+            The value of the 'answer' field if data is a dict,
+            otherwise returns the data unchanged
+        """
+        if isinstance(data, dict):
+            return data.get("answer", None)
+        return data
+
+    @staticmethod
+    def _sort_unordered_list_field(data: Any) -> Any:
+        """
+        Sorts the 'unordered_list' field in a dictionary.
+
+        This enables comparison of unordered lists by converting them to
+        a canonical sorted form.
+
+        Args:
+            data: Dictionary potentially containing an 'unordered_list' field
+
+        Returns:
+            Sorted list if data is a dict with 'unordered_list',
+            otherwise returns data unchanged
+        """
+        if isinstance(data, dict) and "unordered_list" in data:
+            return sorted(data["unordered_list"])
+        return data
 
     @staticmethod
-    def _cast_prediction(reference: dict, prediction: dict) -> None | dict:
+    def _cast_prediction_to_reference_types(
+        reference: dict, prediction: dict
+    ) -> dict | None:
         """
-        Casts the values in the prediction dictionary to match the types
-        of the values in the reference dictionary.
+        Casts prediction values to match reference types.
+
+        Ensures that predictions can be compared with references even when
+        the types differ (e.g., string "123" vs int 123, int 5 vs float 5.0).
+
+        Args:
+            reference: Reference dictionary with expected types
+            prediction: Prediction dictionary to cast
+
+        Returns:
+            Dictionary with casted values, or None if casting fails or
+            prediction is missing required keys
         """
+        if not isinstance(prediction, dict) or not isinstance(reference, dict):
+            return None
+
         casted_prediction = {}
 
         try:
@@ -142,12 +215,12 @@ class TestOfTimeAccuracy(evaluate.Metric):
                 reference_type = type(ref_value)
                 pred_value = prediction[ref_key]
 
-                # Special safeguard: Python allows list("abc") -> ['a', 'b', 'c'].
-                # We don't want to turn strings into character lists.
-                if reference_type == list and not isinstance(pred_value, list):
+                # Safeguard: Python allows list("abc") -> ['a', 'b', 'c']
+                # We don't want to turn strings into character lists
+                if reference_type is list and not isinstance(pred_value, list):
                     return None
 
-                # This handles int("123") -> 123, float(12) -> 12.0, str(100) -> "100"
+                # Cast to reference type: int("123") -> 123, float(12) -> 12.0, etc.
                 casted_prediction[ref_key] = reference_type(pred_value)
 
             return casted_prediction
@@ -156,47 +229,207 @@ class TestOfTimeAccuracy(evaluate.Metric):
             return None
 
     @staticmethod
-    def _fix_age_instruction_label_conflict(pred):
-        if pred is not None:
-            if "age" in pred.keys():
-                # Instruction and label conflict each other in the dataset. Fixes it.
-                pred = {"answer": pred["age"]}
-        return pred
+    def _normalise_list_field_casing(data: dict | None) -> dict | None:
+        """
+        Converts all list items to lowercase for case-insensitive comparison.
+
+        Applied to 'ordered_list' and 'unordered_list' fields to handle variations
+        in capitalization (e.g., "Skating" vs "skating").
+
+        Args:
+            data: Dictionary potentially containing list fields
+
+        Returns:
+            Dictionary with lowercased list items, or None if data is None
+        """
+        if data is None or not isinstance(data, dict):
+            return data
+
+        # Process the first key if it's a list field
+        if data:
+            first_key = next(iter(data.keys()))
+            if first_key in ["ordered_list", "unordered_list"]:
+                data[first_key] = [item.lower() for item in data[first_key]]
+
+        return data
+
+    @staticmethod
+    def _fix_age_field_conflict(prediction: dict | None) -> dict | None:
+        """
+        Fixes a known conflict in the dataset regarding the 'age' field.
+
+        In some dataset samples, the instruction asks for an 'age' field but
+        the reference uses 'answer'. This method normalises the prediction
+        to match the expected format.
+
+        Args:
+            prediction: Prediction dictionary potentially with 'age' field
+
+        Returns:
+            Dictionary with 'age' converted to 'answer', or unchanged if
+            'age' field not present
+        """
+        if prediction is not None and isinstance(prediction, dict):
+            if "age" in prediction:
+                prediction = {"answer": prediction["age"]}
+        return prediction
+
+    def _process_arithmetic_prediction(
+        self, prediction: dict | None, reference: dict | None
+    ) -> tuple[Any, Any]:
+        """
+        Processes a prediction-reference pair for the arithmetic subset.
+
+        Applies arithmetic-specific transformations:
+        1. Fixes age field conflicts
+        2. normalises list casing
+        3. Casts prediction types to match reference
+        4. Sorts unordered lists for comparison
+
+        Args:
+            prediction: Raw prediction dictionary
+            reference: Raw reference dictionary
+
+        Returns:
+            Tuple of (processed_prediction, processed_reference)
+        """
+        prediction = self._fix_age_field_conflict(prediction)
+        prediction = self._normalise_list_field_casing(prediction)
+        reference = self._normalise_list_field_casing(reference)
+        prediction = self._cast_prediction_to_reference_types(reference, prediction)
+
+        # Sort unordered lists for order-independent comparison
+        if reference and "unordered_list" in reference:
+            prediction = self._sort_unordered_list_field(prediction)
+            reference = self._sort_unordered_list_field(reference)
+
+        return prediction, reference
+
+    def _process_semantic_prediction(
+        self, prediction: Any, reference: Any
+    ) -> tuple[str, str]:
+        """
+        Processes a prediction-reference pair for the semantic subset.
+
+        Converts both to strings for comparison since semantic answers
+        may have type mismatches (e.g., int in JSON vs string in reference).
+
+        Args:
+            prediction: Raw prediction value
+            reference: Raw reference value
+
+        Returns:
+            Tuple of (str(prediction), str(reference))
+        """
+        return str(prediction), str(reference)
+
+    def _extract_predictions(
+        self, raw_predictions: list[str], subset: str
+    ) -> list[Any]:
+        """
+        Extracts and preprocesses predictions based on subset type.
+
+        Args:
+            raw_predictions: List of raw prediction strings (e.g., from LLM output)
+            subset: Either 'arithmetic' or 'semantic'
+
+        Returns:
+            List of extracted prediction values
+        """
+        predictions = [self._extract_first_json_object(p) for p in raw_predictions]
+
+        if subset == "semantic":
+            # Since labels are not dicts, we need to extract the value from the LLM's answer field.
+            predictions = [self._extract_answer_field(p) for p in predictions]
+        elif subset == "arithmetic":
+            # Labels and LLMs differ only by the explanation field. Thus, remove.
+            predictions = [self._remove_explanation_field(p) for p in predictions]
+
+        return predictions
+
+    def _extract_references(self, raw_references: list[str], subset: str) -> list[Any]:
+        """
+        Extracts and preprocesses references based on subset type.
+
+        Args:
+            raw_references: List of raw reference strings
+            subset: Either 'arithmetic' or 'semantic'
+
+        Returns:
+            List of extracted reference values
+        """
+        if subset == "arithmetic":
+            # Arithmetic references are Python dict strings that need parsing
+            return [self._parse_reference_label(r) for r in raw_references]
+        else:
+            # Semantic references are used as-is
+            return raw_references
+
+    def _compare_pair(self, prediction: Any, reference: Any, subset: str) -> bool:
+        """
+        Compares a single prediction-reference pair.
+
+        Args:
+            prediction: Processed prediction value
+            reference: Processed reference value
+            subset: Either 'arithmetic' or 'semantic'
+
+        Returns:
+            True if prediction matches reference, False otherwise
+        """
+        if subset == "arithmetic":
+            prediction, reference = self._process_arithmetic_prediction(
+                prediction, reference
+            )
+        elif subset == "semantic":
+            prediction, reference = self._process_semantic_prediction(
+                prediction, reference
+            )
+
+        return prediction == reference
 
     def _compute(
         self,
-        predictions,
-        references,
+        predictions: list[str],
+        references: list[str],
         subset: Literal["arithmetic", "semantic"],
         return_average: bool = True,
-    ):
-        """Returns the scores"""
-        predictions = [self._extract_first_json_object(p) for p in predictions]
-        if subset == "semantic":
-            # Semantic subset's answers are not JSON objects.
-            # Expected answers are always in "answer" field.
-            predictions = [self._get_answer(p) for p in predictions]
-        elif subset == "arithmetic":
-            # Arithmetic subset's answers are JSON objects.
-            # Answer fields vary. Thus, remove explanation field.
-            predictions = [self._pop_explanation(p) for p in predictions]
-            references = [self._parse_label(r) for r in references]
-        else:
-            raise ValueError(f"Invalid subset: {subset}")
-        accuracy = []
-        for pred, ref in zip(predictions, references):
-            if subset == "arithmetic":
-                pred = self._fix_age_instruction_label_conflict(pred)
-                pred = self._cast_prediction(ref, pred)
-                if "unordered_list" in ref:
-                    pred = self._sort_unordered_list(pred)
-                    ref = self._sort_unordered_list(ref)
-            if subset == "semantic":
-                pred = str(pred)
-                ref = str(ref)
-            accuracy.append(
-                pred == ref
-            )  # Semantic subset answer JSON somestimes has int as value. Label is string.
+    ) -> dict[str, float | list[bool]]:
+        """
+        Computes accuracy scores for the Test of Time benchmark.
+
+        Args:
+            predictions: List of prediction strings (LLM outputs)
+            references: List of reference answer strings
+            subset: Benchmark subset - either 'arithmetic' or 'semantic'
+            return_average: If True, returns average accuracy; if False,
+                          returns per-sample correctness
+
+        Returns:
+            Dictionary with 'accuracy' key containing either:
+            - float: average accuracy (if return_average=True)
+            - list[bool]: per-sample correctness (if return_average=False)
+
+        Raises:
+            ValueError: If subset is not 'arithmetic' or 'semantic'
+        """
+        # Validate subset
+        if subset not in ["arithmetic", "semantic"]:
+            raise ValueError(
+                f"Invalid subset: {subset}. Must be 'arithmetic' or 'semantic'."
+            )
+
+        # Extract and preprocess predictions and references
+        predictions = self._extract_predictions(predictions, subset)
+        references = self._extract_references(references, subset)
+
+        # Compare each prediction-reference pair
+        accuracy_scores = [
+            self._compare_pair(pred, ref, subset)
+            for pred, ref in zip(predictions, references)
+        ]
+
+        # Return average or per-sample scores
         if return_average:
-            return {"accuracy": sum(accuracy) / len(accuracy)}
-        return {"accuracy": accuracy}
+            return {"accuracy": sum(accuracy_scores) / len(accuracy_scores)}
+        return {"accuracy": accuracy_scores}

tests/test_arithmetic_scoring.py

@@ -6,6 +6,7 @@ arithmetic_test_cases = {
     "predictions": [
         'JSON = {"explanation": "The war began in 360 BC. Since BC years count backwards, adding 8 years to 360 BC means subtracting 8 from 360, resulting in 352 BC.", "answer": "352 BC"}',
         '```json\n{\n  "explanation": "The dates provided are March 2012, September 2011, June 2017, September 2019, and June 2015. These correspond to visits to Miami, Sydney, Tokyo, London, and Nairobi respectively. The latest date among these is September 2019, which is associated with London. Therefore, London is the last city visited.",\n  "unordered_list": ["Berlin","London"]\n}\n```',
+        '```json\n{\n  "explanation": "The dates provided are March 2012, September 2011, June 2017, September 2019, and June 2015. These correspond to visits to Miami, Sydney, Tokyo, London, and Nairobi respectively. The latest date among these is September 2019, which is associated with London. Therefore, London is the last city visited.",\n  "unordered_list": ["berlin","london"]\n}\n```',
         '```json\n{\n  "explanation": "The dates provided are March 2012, September 2011, June 2017, September 2019, and June 2015. These correspond to visits to Miami, Sydney, Tokyo, London, and Nairobi respectively. The latest date among these is September 2019, which is associated with London. Therefore, London is the last city visited.",\n  "malformed_unordered_list": ["Berlin","London"]\n}\n```',
         ' "To find the date of the second most important game, we need to subtract 7 days from the date of the most important game. We can do this by counting back 7 days from April 14, 2005. April 14 - 7 days = April 7, 2005", "answer": "2005-04-07"}',
         '\n```json\n{\n  "explanation": "Step 1: Determine the time it takes the robot to carry a single box. The robot takes 4 hours, 34 minutes, and 30 seconds to carry 2 boxes. We divide this time by 2 to find the time per box.\\n- Hours: 4 / 2 = 2 hours\\n- Minutes: 34 / 2 = 17 minutes\\n- Seconds: 30 / 2 = 15 seconds\\nSo, it takes the robot 2 hours, 17 minutes, and 15 seconds to carry one box.\\n\\nStep 2: Calculate the total time to carry 25 boxes. We multiply the time per box by the total number of boxes (25).\\n- Total Hours: 2 hours/box * 25 boxes = 50 hours\\n- Total Minutes: 17 minutes/box * 25 boxes = 425 minutes\\n- Total Seconds: 15 seconds/box * 25 boxes = 375 seconds\\n\\nStep 3: Convert the calculated time into the standard H:M:S format by carrying over excess seconds and minutes.\\n- Convert seconds to minutes: 375 seconds is equal to 6 minutes and 15 seconds (since 375 / 60 = 6 with a remainder of 15). We add the 6 minutes to our minutes total.\\n- New total: 50 hours, (425 + 6) minutes, 15 seconds -> 50 hours, 431 minutes, 15 seconds.\\n- Convert minutes to hours: 431 minutes is equal to 7 hours and 11 minutes (since 431 / 60 = 7 with a remainder of 11). We add the 7 hours to our hours total.\\n- New total: (50 + 7) hours, 11 minutes, 15 seconds -> 57 hours, 11 minutes, 15 seconds.\\n\\nThe final time is 57 hours, 11 minutes, and 15 seconds.",\n  "H": 57,\n  "M": 11,\n  "S": 15\n}\n```',
@@ -15,12 +16,13 @@ arithmetic_test_cases = {
         '{"answer": "352 BC"}',
         '{"unordered_list": ["London", "Berlin"]}',
         '{"unordered_list": ["London", "Berlin"]}',
+        '{"unordered_list": ["London", "Berlin"]}',
         '{"answer": "2005-04-07"}',
         '{"H": 57.0, "M": 11.0, "S": 15.0}',
         '{"answer": 3319}',
     ],
-    "result": {"accuracy": 4 / 6},
-    "per_item_accuracy": [True, True, False, False, True, True],
+    "result": {"accuracy": 5 / 7},
+    "per_item_accuracy": [True, True, True,False, False, True, True],
 }
 
 

tests/test_arithmetic_type_casting.py

@@ -16,7 +16,7 @@
 # days, hours, minutes, seconds  <class 'int'>, <class 'int'>, <class 'int'>, <class 'int'>     12
 ##################################################################################################
 
-from test_arithmetic_scoring import TestOfTimeAccuracy
+from test_of_time_accuracy import TestOfTimeAccuracy
 
 
 def test_answer_type_casting():
@@ -29,7 +29,7 @@ def test_answer_type_casting():
         {"answer": "1032"},
     ]
     for ref, pred in zip(references_answer_key, predictions_answer_key):
-        pred_cast = TestOfTimeAccuracy._cast_prediction(ref, pred)
+        pred_cast = TestOfTimeAccuracy._cast_prediction_to_reference_types(ref, pred)
         assert ref == pred_cast
 
 
@@ -37,7 +37,7 @@ def test_unordered_list_type_casting():
     references_unordered_list_key = [{"unordered_list": ["Kyle", "Jason", "Joe"]}]
     predictions_unordered_list_key = [{"unordered_list": ["Kyle", "Jason", "Joe"]}]
     for ref, pred in zip(references_unordered_list_key, predictions_unordered_list_key):
-        pred_cast = TestOfTimeAccuracy._cast_prediction(ref, pred)
+        pred_cast = TestOfTimeAccuracy._cast_prediction_to_reference_types(ref, pred)
         assert ref == pred_cast
 
 
@@ -45,7 +45,7 @@ def test_date_type_casting():
     references_date_key = [{"date": "12/11/2011"}]
     predictions_date_key = [{"date": "12/11/2011"}]
     for ref, pred in zip(references_date_key, predictions_date_key):
-        pred_cast = TestOfTimeAccuracy._cast_prediction(ref, pred)
+        pred_cast = TestOfTimeAccuracy._cast_prediction_to_reference_types(ref, pred)
         assert ref == pred_cast
 
 
@@ -53,7 +53,7 @@ def test_day_time_type_casting():
     references_day_time_keys = [{"day": "+2", "time": "21:44:10"}]
     predictions_day_time_keys = [{"day": "+2", "time": "21:44:10"}]
     for ref, pred in zip(references_day_time_keys, predictions_day_time_keys):
-        pred_cast = TestOfTimeAccuracy._cast_prediction(ref, pred)
+        pred_cast = TestOfTimeAccuracy._cast_prediction_to_reference_types(ref, pred)
         assert ref == pred_cast
 
 
@@ -71,7 +71,7 @@ def test_hms_type_casting():
         {"H": "2.0", "M": "13.0", "S": "30.0"},
     ]
     for ref, pred in zip(references_hms_keys, predictions_hms_keys):
-        pred_cast = TestOfTimeAccuracy._cast_prediction(ref, pred)
+        pred_cast = TestOfTimeAccuracy._cast_prediction_to_reference_types(ref, pred)
         assert ref == pred_cast
 
 
@@ -83,7 +83,7 @@ def test_ordered_list_type_casting():
         {"ordered_list": ["Joe", "Jenny", "Jason", "Dan", "Kyle"]},
     ]
     for ref, pred in zip(references_ordered_list_key, predictions_ordered_list_key):
-        pred_cast = TestOfTimeAccuracy._cast_prediction(ref, pred)
+        pred_cast = TestOfTimeAccuracy._cast_prediction_to_reference_types(ref, pred)
         assert ref == pred_cast
 
 # TODO: Check if I should treat float strings differently, e.g., int(float("18.0"))
@@ -101,7 +101,7 @@ def test_abc_type_casting():
         # {"A": "80.0", "B": "22.0", "C": "20.0"},
     ]
     for ref, pred in zip(references_abc_keys, predictions_abc_keys):
-        pred_cast = TestOfTimeAccuracy._cast_prediction(ref, pred)
+        pred_cast = TestOfTimeAccuracy._cast_prediction_to_reference_types(ref, pred)
         assert ref == pred_cast
 
 
@@ -123,7 +123,7 @@ def test_xyz_type_casting():
         {"X": 2.0, "Y": 4.0, "Z": 36.0},
     ]
     for ref, pred in zip(references_xyz_keys, predictions_xyz_keys):
-        pred_cast = TestOfTimeAccuracy._cast_prediction(ref, pred)
+        pred_cast = TestOfTimeAccuracy._cast_prediction_to_reference_types(ref, pred)
         assert ref == pred_cast
 
 
@@ -143,7 +143,7 @@ def test_hours_minutes_seconds_type_casting():
     for ref, pred in zip(
         references_hours_minutes_seconds_keys, predictions_hours_minutes_seconds_keys
     ):
-        pred_cast = TestOfTimeAccuracy._cast_prediction(ref, pred)
+        pred_cast = TestOfTimeAccuracy._cast_prediction_to_reference_types(ref, pred)
         assert ref == pred_cast
 
 
@@ -161,7 +161,7 @@ def test_hours_minutes_type_casting():
         # {"hours": "5.0", "minutes": "0.0"},
     ]
     for ref, pred in zip(references_hours_minutes_keys, predictions_hours_minutes_keys):
-        pred_cast = TestOfTimeAccuracy._cast_prediction(ref, pred)
+        pred_cast = TestOfTimeAccuracy._cast_prediction_to_reference_types(ref, pred)
         assert ref == pred_cast
 
 
@@ -182,5 +182,5 @@ def test_days_hours_minutes_seconds_type_casting():
         references_days_hours_minutes_seconds_keys,
         predictions_days_hours_minutes_seconds_keys,
     ):
-        pred_cast = TestOfTimeAccuracy._cast_prediction(ref, pred)
+        pred_cast = TestOfTimeAccuracy._cast_prediction_to_reference_types(ref, pred)
         assert ref == pred_cast