Add a chat_template prompt strategy for DPO (#1725)

* Implementing a basic chat_template strategy for DPO datasets

This mimics the sft chat_template strategy such that users can:
* Specify the messages field
* Specify the per message role and content fields
* speicfy the chosen and rejected fields
* Let the tokenizer construct the raw prompt
* Ensure the chosen and rejected fields don't have any prefix tokens

* Adding additional dpo chat template unittests

* Rename test class

src/axolotl/prompt_strategies/dpo/chat_template.py

@@ -0,0 +1,78 @@
+"""
+DPO prompt strategies for using tokenizer chat templates.
+"""
+
+from axolotl.utils.chat_templates import chat_templates
+
+
+def default(
+    cfg, dataset_idx=0, **kwargs
+):  # pylint: disable=possibly-unused-variable,unused-argument
+    ds_cfg = cfg["datasets"][dataset_idx]
+    chat_template_str = chat_templates(cfg.chat_template)
+
+    field_messages = ds_cfg.get("field_messages", "messages")
+    field_chosen = ds_cfg.get("field_chosen", "chosen")
+    field_rejected = ds_cfg.get("field_rejected", "rejected")
+    field_message_role = ds_cfg.get("message_field_role", "role")
+    field_message_content = ds_cfg.get("message_field_content", "content")
+    role_map_inv = ds_cfg.get(
+        "roles",
+        {
+            "user": ["user"],
+            "assistant": ["assistant"],
+            "system": ["system"],
+        },
+    )
+    role_map = {}
+    for target, sources in role_map_inv.items():
+        for source in sources:
+            role_map[source] = target
+
+    def transform_fn(sample, tokenizer=None):
+        messages = sample[field_messages]
+        messages = [
+            {
+                "role": role_map[m[field_message_role]],
+                "content": m[field_message_content],
+            }
+            for m in messages
+        ]
+        chosen = {
+            "role": role_map[sample[field_chosen][field_message_role]],
+            "content": sample[field_chosen][field_message_content],
+        }
+        rejected = {
+            "role": role_map[sample[field_rejected][field_message_role]],
+            "content": sample[field_rejected][field_message_content],
+        }
+
+        result = {}
+        result["prompt"] = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            chat_template=chat_template_str,
+            tokenize=False,
+        )
+
+        result["chosen"] = tokenizer.apply_chat_template(
+            [chosen],
+            add_generation_prompt=False,
+            chat_template=chat_template_str,
+            tokenize=False,
+        )
+        chosen_strip_index = result["chosen"].find(chosen["content"])
+        result["chosen"] = result["chosen"][chosen_strip_index:]
+
+        result["rejected"] = tokenizer.apply_chat_template(
+            [rejected],
+            add_generation_prompt=False,
+            chat_template=chat_template_str,
+            tokenize=False,
+        )
+        rejected_strip_index = result["rejected"].find(rejected["content"])
+        result["rejected"] = result["rejected"][rejected_strip_index:]
+
+        return result
+
+    return transform_fn

src/axolotl/utils/data/rl.py

@@ -1,4 +1,5 @@
 """data handling specific to DPO"""
+
 import inspect
 import logging
 from functools import partial

src/axolotl/utils/tokenization.py

@@ -62,7 +62,7 @@ def process_tokens_for_rl_debug(tokens, color, tokenizer, text_only):
     """Helper function to process and color tokens."""
     colored_tokens = [
         color_token_for_rl_debug(tokenizer.decode(token), token, color, text_only)
-        for token in tokenizer.encode(tokens)
+        for token in tokenizer.encode(tokens, add_special_tokens=False)
     ]
     return colored_tokens