axolotl-ai-cloud · Feb 29, 2024
Showing with 19 additions and 0 deletions.

+19 −0 src/axolotl/prompt_strategies/dpo/chatml.py
diff --git a/src/axolotl/prompt_strategies/dpo/chatml.py b/src/axolotl/prompt_strategies/dpo/chatml.py
@@ -23,6 +23,25 @@ def transform_fn(sample):
 
     return transform_fn
 
+def llamafactory(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
+    def transform_fn(sample):
+        data = {}
+        if "instruction" in sample and sample["instruction"]:
+            data["prompt"] = (
+                f"<|im_start|>system\n{sample['instruction']}<|im_end|>\n"
+                f"<|im_start|>user\n{sample['input']}<|im_end|>\n<|im_start|>assistant\n"
+            )
+        else:
+            data["prompt"] = (
+                f"<|im_start|>user\n{sample['input']}<|im_end|>\n<|im_start|>assistant\n"
+            )
+
+        data["chosen"] = sample["output"][0]
+        data["rejected"] = sample["output"][1]
+
+        return data
+
+    return transform_fn
 
 def icr(
     cfg,