Add dataset_output_field_values_to_texts option

Author: yuiseki

recipes/A5000_24GB_x8/fake-news-detector-en-1.5T.yaml

@@ -5,6 +5,9 @@ output_base_dir: /data/output
 dataset_id: mrm8488/fake-news
 dataset_input_field_name: text
 dataset_output_field_name: label
+dataset_output_field_values_to_texts:
+  0: "Real"
+  1: "Fake"
 dataset_train_split_seed: 42
 dataset_train_split_test_size: 0.2
 lora_r: 8
@@ -20,16 +23,16 @@ inference_max_new_tokens: 2
 evaluations:
   -
     prompt: "Donald Trump has never been President of the United States."
-    expected_output: "1"
+    expected_output: "Fake"
   -
     prompt: "The Earth is flat."
-    expected_output: "1"
+    expected_output: "Fake"
   -
     prompt: "Martians visited Japan in 2011."
-    expected_output: "1"
+    expected_output: "Fake"
   -
     prompt: "The World Trade Center collapsed when the plane hit it."
-    expected_output: "0"
+    expected_output: "Real"
   -
-    expected_output: "0"
     prompt: "The United States is a country in North America."
+    expected_output: "Real"

recipes/A5000_24GB_x8/fake-news-detector-en.yaml

@@ -5,6 +5,9 @@ output_base_dir: /data/output
 dataset_id: mrm8488/fake-news
 dataset_input_field_name: text
 dataset_output_field_name: label
+dataset_output_field_values_to_texts:
+  0: "Real"
+  1: "Fake"
 dataset_train_split_seed: 42
 dataset_train_split_test_size: 0.2
 lora_r: 8
@@ -20,16 +23,16 @@ inference_max_new_tokens: 2
 evaluations:
   -
     prompt: "Donald Trump has never been President of the United States."
-    expected_output: "1"
+    expected_output: "Fake"
   -
     prompt: "The Earth is flat."
-    expected_output: "1"
+    expected_output: "Fake"
   -
     prompt: "Martians visited Japan in 2011."
-    expected_output: "1"
+    expected_output: "Fake"
   -
     prompt: "The World Trade Center collapsed when the plane hit it."
-    expected_output: "0"
+    expected_output: "Real"
   -
-    expected_output: "0"
     prompt: "The United States is a country in North America."
+    expected_output: "Real"

src/train.py

@@ -125,6 +125,11 @@ def prepare_train_data(dataset_id):
         data_df["text"] = data_df[input_field_name].apply(lambda x: simple_template_for_pretrain(x))
     else:
         output_field_name = train_config["dataset_output_field_name"]
+        if "dataset_output_field_values_to_texts" in train_config:
+            output_field_values_to_texts = train_config["dataset_output_field_values_to_texts"]
+            data_df[output_field_name] = data_df[output_field_name].apply(
+                lambda x: output_field_values_to_texts.get(x, x)
+            )
         if "dataset_context_field_name" in train_config:
             context_field_name = train_config["dataset_context_field_name"]
             if "dataset_context_hint" not in train_config: