fix tests

Signed-off-by: Sukriti-Sharma4 <[email protected]>
foundation-model-stack · May 24, 2024 · 1f6bb04 · 1f6bb04
1 parent 45827ce
commit 1f6bb04
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/tests/utils/test_data_utils.py b/tests/utils/test_data_utils.py
@@ -31,7 +31,8 @@ def test_formatting_function():
     template = "### Input: {{Tweet text}} \n\n ### Response: {{text_label}}"
     # First response from the data file that is read.
     expected_response = (
-        "### Input: No this is my first job \n\n ### Response: no complaint"
+        "### Input: @HMRCcustomers No this is my first job"
+        + " \n\n ### Response: no complaint"
     )
     formatted_dataset, dataset_text_field = data_utils.formatting_function(
         json_dataset, template
@@ -46,7 +47,8 @@ def test_formatting_function_adds_eos_token():
     template = "### Input: {{Tweet text}} \n\n ### Response: {{text_label}}"
     # First response from the data file that is read.
     expected_response = (
-        "### Input: No this is my first job \n\n ### Response: no complaintEOS"
+        "### Input: @HMRCcustomers No this is my first job"
+        + " \n\n ### Response: no complaintEOS"
     )
     formatted_dataset, dataset_text_field = data_utils.formatting_function(
         json_dataset, template, "EOS"