Add HFSummaryPairs class & fix AnthropicRLHF parsing
diff --git a/model/model_training/configs/config_rm.yaml b/model/model_training/configs/config_rm.yaml
index 447b75acca..1b66a1bb70 100644
--- a/model/model_training/configs/config_rm.yaml
+++ b/model/model_training/configs/config_rm.yaml
@@ -85,7 +85,7 @@ oasst-rm-1-pythia-1.4b:
     - webgpt:
         val_split: 0.05
         max_val_set: 1000
-    - hf_summary:
+    - hf_summary_pairs:
         fraction: 0.1
         max_val_set: 250
   use_custom_sampler: true
@@ -97,10 +97,10 @@ oasst-rm-1-pythia-1.4b:
   dtype: float32
   max_length: 2048
   use_flash_attention: true
-  warmup_steps: 10
-  gradient_accumulation_steps: 2
-  per_device_train_batch_size: 2
-  per_device_eval_batch_size: 6
+  warmup_steps: 50
+  gradient_accumulation_steps: 4
+  per_device_train_batch_size: 1
+  per_device_eval_batch_size: 5
   num_train_epochs: 2
   eval_steps: 500
   save_steps: 1000
diff --git a/model/model_training/custom_datasets/__init__.py b/model/model_training/custom_datasets/__init__.py
index 01b415d4ec..2b7e085442 100644
--- a/model/model_training/custom_datasets/__init__.py
+++ b/model/model_training/custom_datasets/__init__.py
@@ -19,7 +19,7 @@
     WebGPT,
 )
 from model_training.custom_datasets.rank_datasets import AugmentedOA
-from model_training.custom_datasets.summarization import HFSummary, SummarizationDataset
+from model_training.custom_datasets.summarization import HFSummary, HFSummaryPairs, SummarizationDataset
 from model_training.custom_datasets.toxic_conversation import ProsocialDialogue, ProsocialDialogueExplaination
 from model_training.custom_datasets.translation import WMT2019, DiveMT, TEDTalk
 from sklearn.model_selection import train_test_split
@@ -50,6 +50,7 @@
     "private_tuning",
     "alpaca",
     "hf_summary",
+    "hf_summary_pairs",
 ]
 
 RM_DATASETS = [
@@ -57,6 +58,7 @@
     "augment_oasst",
     "anthropic_rlhf",
     "hf_summary",
+    "hf_summary_pairs",
     "shp",
     "hellaswag",
     "webgpt",
@@ -140,6 +142,9 @@ def get_one_dataset(
     elif dataset_name == "hf_summary":
         train = HFSummary(split="train", mode=mode)
         eval = HFSummary(split="valid1", mode=mode)
+    elif dataset_name == "hf_summary_pairs":
+        train = HFSummaryPairs(split="train", mode=mode)
+        eval = HFSummaryPairs(split="valid1", mode=mode)
     elif dataset_name == "augment_oasst":
         # reward model mode only
         assert mode == "rm"
diff --git a/model/model_training/custom_datasets/rank_datasets.py b/model/model_training/custom_datasets/rank_datasets.py
index 54a0510e29..a45c7f237e 100644
--- a/model/model_training/custom_datasets/rank_datasets.py
+++ b/model/model_training/custom_datasets/rank_datasets.py
@@ -183,10 +183,15 @@ class AnthropicRLHF(Dataset):
     name = "anthropic_rlhf"
 
     @staticmethod
-    def _split_dialogue(text: str):
+    def _split_dialogue(text: str) -> list[tuple[str, str]]:
         lines = text.split("\n\n")
 
-        dialogue = []
+        dialogue: list[tuple[str, str]] = []
+
+        # go over messages and combine consecutive messages from the
+        # same speaker (OA v1 expects alternating roles)
+        role = None
+        messages = []
         for line in lines:
             if line.startswith("Human:"):
                 speaker = "Human"
@@ -196,16 +201,25 @@ def _split_dialogue(text: str):
                 message = line[11:]
             else:
                 continue
-            dialogue.append((speaker, message.strip()))
+            if role != speaker:
+                if role is not None:
+                    dialogue.append((role, "\n".join(messages)))
+                    messages = []
+                role = speaker
+            messages.append(message.strip())
+
+        if role is not None and len(messages) > 0:
+            dialogue.append((role, "\n".join(messages)))
 
         return dialogue
 
-    def __init__(self, split="train") -> None:
+    def __init__(self, split: str = "train") -> None:
         super().__init__()
         assert split in ("train", "test")
         self.split = split
         self.data = []
         dataset = load_dataset("Anthropic/hh-rlhf")[split]
+
         for entry in dataset:
             chosen = entry["chosen"]
 
@@ -215,14 +229,17 @@ def __init__(self, split="train") -> None:
             rejected = entry["rejected"]
             chosen = self._split_dialogue(chosen)
             rejected = self._split_dialogue(rejected)
+            assert rejected[0][0] == "Human" and chosen[0][0] == "Human"
 
-            prefix = [line for (speaker, line) in chosen[:-1]]
-            good_reply = chosen[-1][1]  # last part of dialog, the text
-            bad_reply = rejected[-1][1]  # last part of dialog, the text
-            self.data.append((prefix, [good_reply, bad_reply]))
+            # only very few items have non matching lengths
+            if len(rejected) == len(chosen):
+                prefix = [line for (speaker, line) in chosen[:-1]]
+                good_reply = chosen[-1][1]  # last part of dialog, the text
+                bad_reply = rejected[-1][1]  # last part of dialog, the text
+                self.data.append((prefix, [good_reply, bad_reply]))
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.data)
 
-    def __getitem__(self, index):
+    def __getitem__(self, index: int) -> tuple[str, list[str]]:
         return self.data[index]
diff --git a/model/model_training/custom_datasets/summarization.py b/model/model_training/custom_datasets/summarization.py
index e681110bc4..d9a7efdfe0 100644
--- a/model/model_training/custom_datasets/summarization.py
+++ b/model/model_training/custom_datasets/summarization.py
@@ -80,6 +80,89 @@ def __getitem__(self, idx):
         return (context, summary)
 
 
+SUMMARIZATION_PROMPTS = [
+    "Please summarize the following content:\n{}",
+    "Write me a summary for the following article:\n{}",
+    "Kindly sum up the following information: {}",
+    "Please summarize the following text for me:\n{}",
+    "Give me a summary of the following text:\n\n{}",
+    "Describe the following information in brief: {}",
+    "Will you kindly summarize the following paragraph for me?\n{}",
+    "Summarize this: {}",
+    "TLDR this: {}",
+    "{}\nTLDR;",
+    "{}\n\nTL;DR",
+    "{} tl;dr:",
+    "{}\nPlease summarize the content above",
+    "{} Please summarize the preceding statements.",
+]
+
+
+class HFSummaryPairs(Dataset):
+    """
+    Simplified version of the HFSummary class which uses the original examples
+    of the OpenAI dataset.
+    https://huggingface.co/datasets/openai/summarize_from_feedback
+    """
+
+    def __init__(self, split="train", mode="sft", conf_threshold=-1) -> None:
+        super().__init__()
+        assert split in ("train", "valid1", "valid2", "test")
+        assert mode in ("sft", "rm", "rl")
+        self.mode = mode
+
+        self.posts = []
+        self.summary_pairs = []
+
+        major_split = split if "train" == split else "validation"
+        dataset = load_dataset("openai/summarize_from_feedback", "comparisons")[major_split]
+        for data in dataset:
+            if (
+                "extra" in data
+                and "confidence" in data["extra"]
+                and data["extra"]["confidence"] is not None
+                and conf_threshold > data["extra"]["confidence"]
+            ):
+                print("skipping {}".format(data["info"]["id"]))
+                continue
+
+            if split != "train" and split != data["split"]:
+                continue
+
+            if "article" in data["info"] and data["info"]["article"] is not None:
+                context = data["info"]["article"]
+            elif "post" in data["info"]:
+                context = data["info"]["post"]
+
+            self.posts.append(context)
+            pos, neg = (0, 1) if data["choice"] == 0 else (1, 0)
+            self.summary_pairs.append((data["summaries"][pos]["text"].strip(), data["summaries"][neg]["text"].strip()))
+
+    def __len__(self) -> int:
+        return len(self.posts)
+
+    def __getitem__(self, index: int) -> tuple | list:
+        if index < 0 or index >= len(self.posts):
+            raise IndexError()
+
+        context = self.posts[index]
+        # return pairs of comparison
+        good_summary, bad_summary = self.summary_pairs[index]
+        prompt = random.choice(SUMMARIZATION_PROMPTS)
+
+        # pair very big
+        # we are going to do some sampling
+        # not optimal but good for now
+        if self.mode == "sft":
+            return [prompt.format(context), good_summary]
+        elif self.mode == "rl":
+            return (prompt.format(context),)
+        elif self.mode == "rm":
+            return [prompt.format(context)], [good_summary, bad_summary]
+
+        raise RuntimeError(f"Unsupported mode '{self.mode}'")
+
+
 class HFSummary(Dataset):
     """
     Human feedback data from OpenAI
@@ -90,13 +173,6 @@ class HFSummary(Dataset):
 
     """
 
-    PROMPTS = [
-        "Please summarize the following content:\n{}",
-        "{}\nTLDR;",
-        "{}\nPlease summarize the content above",
-        "Write a summary for the following article:\n{}",
-    ]
-
     def __init__(self, split="train", mode="sft", conf_threshold=-1, max_comparison_per_sample=5) -> None:
         super().__init__()
         assert split in ("train", "valid1", "valid2", "test")
@@ -165,17 +241,17 @@ def get_sorted_ranks(comparison_pairs):
 
         return sorted_elements
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.index2summary)
 
-    def __getitem__(self, index):
-        if index >= len(self.index2summary):
+    def __getitem__(self, index) -> tuple | list:
+        if index < 0 or index >= len(self.index2summary):
             raise IndexError()
 
-        context = self.index2summary.get(index)
+        context = self.index2summary[index]
         # return pairs of comparison
         rows = self.summaries[context]
-        prompt = random.choice(self.PROMPTS)
+        prompt = random.choice(SUMMARIZATION_PROMPTS)
 
         # pair very big
         # we are going to do some sampling
diff --git a/model/model_training/tests/test_ranking_collator.py b/model/model_training/tests/test_ranking_collator.py
index 363beb0be8..fb8938e6f8 100644
--- a/model/model_training/tests/test_ranking_collator.py
+++ b/model/model_training/tests/test_ranking_collator.py
@@ -12,17 +12,17 @@ def test_rm_datasets():
     # dummy configuration
     config = Namespace(cache_dir=".cache", model_name="EleutherAI/pythia-70m-deduped")
 
-    dataset_names = ["webgpt", "hf_summary", "hellaswag", "shp", "anthropic_rlhf"]
+    dataset_names = ["anthropic_rlhf", "hf_summary_pairs", "webgpt", "hellaswag", "shp", "hf_summary"]
     for name in dataset_names:
         train, val = get_one_dataset(conf=config, dataset_name=name, mode="rm")
-        print(f"dataset: {name}  (train ({type(train)}): {len(train)}, val({type(val)}): {len(val)})")
+        print(f"dataset: '{name}' (train ({type(train)}): {len(train)}, val({type(val)}): {len(val)})")
 
         avg_number_continuations = sum(len(x[1]) for x in train) / len(train)
         num_more_than_two = sum(1 if len(x[1]) > 2 else 0 for x in train)
         print(f"Average number of continuations: {avg_number_continuations} (with >2: {num_more_than_two})")
 
-        for i in range(2):
-            item = train[i]
+        for i in range(10):
+            item = train[i + 100]
             print(f"[{i}] Prefix: {item[0]}")
             continuations = item[1]
             print(f"[{i}] Continuations ({len(continuations)}):")

- add simpler version of HFSummaries which directly uses the pairs of the OpenAI dataset
- combine consecutive messages of same speaker during loading of AnthropicRLHF dataset, ignore examples with different structure between positive & negative example (only very few)