From adab6c516a8e36268c03e373549a0ea952669495 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Tue, 12 Mar 2024 17:17:10 -0400
Subject: [PATCH] add handling for argilla dpo-mix (#1397)

---
 src/axolotl/prompt_strategies/dpo/chatml.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/src/axolotl/prompt_strategies/dpo/chatml.py b/src/axolotl/prompt_strategies/dpo/chatml.py
index e8c7f4088c..585696e29a 100644
--- a/src/axolotl/prompt_strategies/dpo/chatml.py
+++ b/src/axolotl/prompt_strategies/dpo/chatml.py
@@ -24,6 +24,25 @@ def transform_fn(sample):
     return transform_fn
 
 
+def argilla_chat(
+    cfg,
+    **kwargs,
+):  # pylint: disable=possibly-unused-variable,unused-argument
+    """
+    for argilla/dpo-mix-7k conversations
+    """
+
+    def transform_fn(sample):
+        sample[
+            "prompt"
+        ] = f"<|im_start|>user\n{sample['chosen'][0]['content']}<|im_end|>\n<|im_start|>assistant\n"
+        sample["chosen"] = f"{sample['chosen'][1]['content']}<|im_end|>"
+        sample["rejected"] = f"{sample['rejected'][1]['content']}<|im_end|>"
+        return sample
+
+    return transform_fn
+
+
 def icr(
     cfg,
     **kwargs,