Skip to content

Commit e194caa

Browse files
authored
Fix remap_column_names (#140)
When I try to do the following, I got error: ```python ds = Dataset.from_dict( { "question": ["question"], "answer": ["answer"], "contexts": [["context"]], } ) from ragas import evaluate from ragas.metrics import Faithfulness evaluate(dataset =ds, metrics=[Faithfulness(batch_size=1)]) ``` ``` KeyError: "Column ground_truths not in the dataset. Current columns in the dataset: ['question', 'answer', 'contexts']" ``` But `ground_truths ` is not needed for `Faithfulness` . This PR is to fix it.
1 parent d0bcea1 commit e194caa

File tree

2 files changed

+21
-0
lines changed

2 files changed

+21
-0
lines changed

src/ragas/validation.py

+1
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ def remap_column_names(dataset: Dataset, column_map: dict[str, str]) -> Dataset:
99
"""
1010
Remap the column names in case dataset uses different column names
1111
"""
12+
column_map = {k: v for k, v in column_map.items() if v in dataset.column_names}
1213
inverse_column_map = {v: k for k, v in column_map.items()}
1314
return dataset.from_dict(
1415
{inverse_column_map[name]: dataset[name] for name in column_map.values()}

tests/unit/test_validation.py

+20
Original file line numberDiff line numberDiff line change
@@ -105,3 +105,23 @@ def test_column_remap(column_map):
105105
remapped_dataset = remap_column_names(TEST_DATASET, column_map)
106106

107107
assert remapped_dataset.column_names == list(column_map.keys())
108+
109+
110+
def test_column_remap_omit():
111+
TEST_DATASET = Dataset.from_dict(
112+
{
113+
"query": [""],
114+
"answer": [""],
115+
"contexts": [[""]],
116+
}
117+
)
118+
119+
column_map = {
120+
"question": "query",
121+
"contexts": "contexts",
122+
"answer": "answer",
123+
"ground_truths": "ground_truths",
124+
}
125+
126+
remapped_dataset = remap_column_names(TEST_DATASET, column_map)
127+
assert remapped_dataset.column_names == ["question", "contexts", "answer"]

0 commit comments

Comments
 (0)