From cdde04c38bfb0848fdf5e5fb49288a982ecbfb13 Mon Sep 17 00:00:00 2001
From: nopperl <54780682+nopperl@users.noreply.github.com>
Date: Fri, 2 Feb 2024 13:25:45 +0100
Subject: [PATCH] support user-defined prompt processing strategies for dpo

---
 src/axolotl/prompt_strategies/dpo/__init__.py |  5 ++-
 src/axolotl/prompt_strategies/dpo/chatml.py   |  8 +++--
 .../prompt_strategies/dpo/user_defined.py     | 31 +++++++++++++++++++
 src/axolotl/prompt_strategies/dpo/zephyr.py   |  2 +-
 src/axolotl/utils/data.py                     |  2 +-
 5 files changed, 40 insertions(+), 8 deletions(-)
 create mode 100644 src/axolotl/prompt_strategies/dpo/user_defined.py

diff --git a/src/axolotl/prompt_strategies/dpo/__init__.py b/src/axolotl/prompt_strategies/dpo/__init__.py
index 3c1c808005..8bd430f912 100644
--- a/src/axolotl/prompt_strategies/dpo/__init__.py
+++ b/src/axolotl/prompt_strategies/dpo/__init__.py
@@ -8,14 +8,13 @@
 LOG = logging.getLogger("axolotl")
 
 
-def load(strategy, cfg):
+def load(strategy, cfg, **kwargs):
     try:
         load_fn = strategy.split(".")[-1]
         strategy = ".".join(strategy.split(".")[:-1])
         mod = importlib.import_module(f".{strategy}", "axolotl.prompt_strategies.dpo")
         func = getattr(mod, load_fn)
-        load_kwargs = {}
-        return func(cfg, **load_kwargs)
+        return func(cfg, **kwargs)
     except Exception:  # pylint: disable=broad-exception-caught
         LOG.warning(f"unable to load strategy {strategy}")
         return None
diff --git a/src/axolotl/prompt_strategies/dpo/chatml.py b/src/axolotl/prompt_strategies/dpo/chatml.py
index 8f62a5088e..0d19d65cbb 100644
--- a/src/axolotl/prompt_strategies/dpo/chatml.py
+++ b/src/axolotl/prompt_strategies/dpo/chatml.py
@@ -5,6 +5,7 @@
 
 def argilla(
     cfg,
+    **kwargs,
 ):  # pylint: disable=possibly-unused-variable,unused-argument
     def transform_fn(sample):
         if "system" in sample and sample["system"]:
@@ -25,6 +26,7 @@ def transform_fn(sample):
 
 def icr(
     cfg,
+    **kwargs,
 ):  # pylint: disable=possibly-unused-variable,unused-argument
     """
     chatml transforms for datasets with system, input, chosen, rejected
@@ -48,7 +50,7 @@ def transform_fn(sample):
     return transform_fn
 
 
-def intel(cfg):  # pylint: disable=possibly-unused-variable,unused-argument
+def intel(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
     """
     For Intel Orca DPO Pairs
     """
@@ -70,7 +72,7 @@ def transform_fn(sample):
     return transform_fn
 
 
-def prompt_pairs(cfg):  # pylint: disable=possibly-unused-variable,unused-argument
+def prompt_pairs(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
     def transform_fn(sample):
         if "system" in sample and sample["system"]:
             sample["prompt"] = (
@@ -88,7 +90,7 @@ def transform_fn(sample):
     return transform_fn
 
 
-def ultra(cfg):  # pylint: disable=possibly-unused-variable,unused-argument
+def ultra(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
     """
     for ultrafeedback binarized conversations
     """
diff --git a/src/axolotl/prompt_strategies/dpo/user_defined.py b/src/axolotl/prompt_strategies/dpo/user_defined.py
new file mode 100644
index 0000000000..5ba488ca3d
--- /dev/null
+++ b/src/axolotl/prompt_strategies/dpo/user_defined.py
@@ -0,0 +1,31 @@
+"""
+User-defined DPO strategies
+"""
+
+
+def default(cfg, dataset_idx=0, **kwargs):
+    ds_cfg = cfg["datasets"][dataset_idx]
+    field_prompt = ds_cfg.get("field_prompt", "prompt")
+    field_system = ds_cfg.get("field_system", "system")
+    field_chosen = ds_cfg.get("field_chosen", "chosen")
+    field_rejected = ds_cfg.get("field_rejected", "rejected")
+    prompt_format = ds_cfg.get("prompt_format")
+    if not prompt_format:
+        prompt_format = "{" + field_prompt + "}"
+    chosen_format = ds_cfg.get("chosen_format")
+    if not chosen_format:
+        chosen_format = "{" + field_chosen + "}"
+    rejected_format = ds_cfg.get("rejected_format")
+    if not rejected_format:
+        rejected_format = "{" + field_rejected + "}"
+
+    def transform_fn(sample):
+        if "{" + field_system + "}" in prompt_format and field_system in sample and sample[field_system]:
+            sample["prompt"] = prompt_format.format(system=sample[field_system], prompt=sample[field_prompt])
+        else:
+            sample["prompt"] = prompt_format.format(prompt=sample["prompt"])
+        sample["chosen"] = chosen_format.format(chosen=sample[field_chosen])
+        sample["rejected"] = rejected_format.format(rejected=sample[field_rejected])
+        return sample
+
+    return transform_fn
diff --git a/src/axolotl/prompt_strategies/dpo/zephyr.py b/src/axolotl/prompt_strategies/dpo/zephyr.py
index 02bce8a338..9eb8950091 100644
--- a/src/axolotl/prompt_strategies/dpo/zephyr.py
+++ b/src/axolotl/prompt_strategies/dpo/zephyr.py
@@ -3,7 +3,7 @@
 """
 
 
-def nectar(cfg):  # pylint: disable=possibly-unused-variable,unused-argument
+def nectar(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
     def transform_fn(sample):
         data = {}
         data["prompt"] = (
diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py
index 5e6ceb6cbe..e43e1ef706 100644
--- a/src/axolotl/utils/data.py
+++ b/src/axolotl/utils/data.py
@@ -914,7 +914,7 @@ def load_split(dataset_cfgs, _cfg):
         for i, data_set in enumerate(split_datasets):
             _type = dataset_cfgs[i]["type"]
             if _type:
-                ds_transform_fn = load_dpo(_type, _cfg)
+                ds_transform_fn = load_dpo(_type, _cfg, dataset_idx=i)
                 split_datasets[i] = data_set.map(
                     ds_transform_fn,
                     desc="Mapping RL Dataset",