diff --git a/src/evaluate/evaluator/utils.py b/src/evaluate/evaluator/utils.py index e364276d..8fc302c7 100644 --- a/src/evaluate/evaluator/utils.py +++ b/src/evaluate/evaluator/utils.py @@ -8,6 +8,9 @@ def __init__(self, dataset: Dataset, key: str): self.dataset = dataset self.key = key + def __repr__(self): + return f"DatasetColumn(key={self.key!r}, len={len(self)})" + def __len__(self): return len(self.dataset) @@ -65,6 +68,16 @@ def __init__( self.first_key = first_key self.second_key = second_key + def __repr__(self): + return ( + f"DatasetColumnPair(" + f"first_col={self.first_col!r}, " + f"second_col={self.second_col!r}, " + f"first_key={self.first_key!r}, " + f"second_key={self.second_key!r}, " + f"len={len(self)})" + ) + def __len__(self): return len(self.dataset) diff --git a/tests/test_evaluator_utils.py b/tests/test_evaluator_utils.py new file mode 100644 index 00000000..fcbf7331 --- /dev/null +++ b/tests/test_evaluator_utils.py @@ -0,0 +1,20 @@ +from datasets import Dataset + +from evaluate.evaluator.utils import DatasetColumn, DatasetColumnPair + + +def test_dataset_column_repr(): + dataset = Dataset.from_dict({"text": ["a", "b", "c"]}) + column = DatasetColumn(dataset, "text") + + assert repr(column) == "DatasetColumn(key='text', len=3)" + + +def test_dataset_column_pair_repr(): + dataset = Dataset.from_dict({"text": ["a", "b", "c"], "label": [0, 1, 0]}) + column_pair = DatasetColumnPair(dataset, "text", "label", "prediction", "reference") + + assert repr(column_pair) == ( + "DatasetColumnPair(first_col='text', second_col='label', " + "first_key='prediction', second_key='reference', len=3)" + )