Skip to content

Commit 96872f5

Browse files
FIX: Preserve text columns when loading human-labeled CSVs (#2980)
Co-authored-by: hannahwestra25 <hannahwestra@microsoft.com>
1 parent fa829c3 commit 96872f5

2 files changed

Lines changed: 45 additions & 2 deletions

File tree

‎pyrit/score/scorer_evaluation/human_labeled_dataset.py‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
from pyrit.score.scorer_evaluation.metrics_type import MetricsType
1616

1717
if TYPE_CHECKING:
18+
from collections.abc import Hashable
19+
1820
from pyrit.models.harm_definition import HarmDefinition
1921

2022
logger = logging.getLogger(__name__)
@@ -271,11 +273,21 @@ def from_csv(
271273
if not harm_definition_version and parsed_harm_definition_version:
272274
harm_definition_version = parsed_harm_definition_version
273275

276+
# Preserve text before pandas can infer numeric or boolean values.
277+
text_dtypes: dict[Hashable, type[str]] = dict.fromkeys(
278+
(
279+
cls.STANDARD_ASSISTANT_RESPONSE_COL,
280+
cls.STANDARD_OBJECTIVE_COL,
281+
cls.STANDARD_HARM_COL,
282+
cls.STANDARD_DATA_TYPE_COL,
283+
),
284+
str,
285+
)
274286
# Try UTF-8 first, fall back to latin-1 for files with special characters
275287
try:
276-
eval_df = pd.read_csv(csv_path, comment="#", encoding="utf-8")
288+
eval_df = pd.read_csv(csv_path, comment="#", encoding="utf-8", dtype=text_dtypes)
277289
except UnicodeDecodeError:
278-
eval_df = pd.read_csv(csv_path, comment="#", encoding="latin-1")
290+
eval_df = pd.read_csv(csv_path, comment="#", encoding="latin-1", dtype=text_dtypes)
279291

280292
# Drop rows where every column is NaN (e.g. trailing blank lines in the CSV)
281293
eval_df = eval_df.dropna(how="all")

‎tests/unit/score/test_human_labeled_dataset.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
# Copyright (c) Microsoft Corporation.
22
# Licensed under the MIT license.
33

4+
from pathlib import Path
5+
46
import pytest
57

68
from pyrit.common.path import SCORER_EVALS_PATH
@@ -13,6 +15,35 @@
1315
)
1416

1517

18+
@pytest.mark.parametrize("value", ["00123", "1e3", "1.2300", "True"])
19+
@pytest.mark.parametrize("encoding", ["utf-8", "latin-1"])
20+
@pytest.mark.parametrize("metrics_type", [MetricsType.OBJECTIVE, MetricsType.HARM])
21+
def test_from_csv_preserves_numeric_and_boolean_text(
22+
tmp_path: Path, value: str, encoding: str, metrics_type: MetricsType
23+
) -> None:
24+
column = "objective" if metrics_type == MetricsType.OBJECTIVE else "harm_category"
25+
label = "1" if metrics_type == MetricsType.OBJECTIVE else "0.5"
26+
csv_path = tmp_path / "text.csv"
27+
csv_path.write_text(
28+
f"assistant_response,{column},human_score,notes\n{value},00456,{label},café\n",
29+
encoding=encoding,
30+
)
31+
32+
dataset = HumanLabeledDataset.from_csv(csv_path=csv_path, metrics_type=metrics_type, version="1.0")
33+
entry = dataset.entries[0]
34+
piece = entry.conversation[0].message_pieces[0]
35+
36+
assert piece.original_value == value
37+
assert piece.converted_value == value
38+
if isinstance(entry, ObjectiveHumanLabeledEntry):
39+
assert entry.objective == "00456"
40+
assert entry.human_scores == [True]
41+
else:
42+
assert isinstance(entry, HarmHumanLabeledEntry)
43+
assert entry.harm_category == "00456"
44+
assert entry.human_scores == [0.5]
45+
46+
1647
@pytest.fixture
1748
def sample_messages():
1849
return [

0 commit comments

Comments
 (0)