update for data format

Author: yihedeng9

spin/alignment/data.py

@@ -31,9 +31,9 @@ def _strip_prefix(s, pattern):
         rejected_messages = example["generated"][1:]
         example["text_chosen"] = tokenizer.apply_chat_template(chosen_messages, tokenize=False)
         example["text_rejected"] = tokenizer.apply_chat_template(rejected_messages, tokenize=False)
-        # example["text_prompt"] = tokenizer.apply_chat_template(
-        #     prompt_messages, tokenize=False, add_generation_prompt=True
-        # )
+        example["text_prompt"] = tokenizer.apply_chat_template(
+            prompt_messages, tokenize=False, add_generation_prompt=True
+        )
         example["text_chosen"] = _strip_prefix(example["text_chosen"], assistant_prefix)
         example["text_rejected"] = _strip_prefix(example["text_rejected"], assistant_prefix)
     else:

spin/run_spin.py

@@ -89,7 +89,7 @@ def main():
     # Replace column names with what TRL needs, text_chosen -> chosen and text_rejected -> rejected
     for split in ["train", "test"]:
         raw_datasets[split] = raw_datasets[split].rename_columns(
-            {"text_real": "real", "text_generated": "generated"}
+            {"text_prompt": "prompt", "text_real": "real", "text_generated": "generated"}
         )
 
     torch_dtype = (