Update Common Pile Formatting

src/configs/common_pile_ultra_permissive.yaml

@@ -31,6 +31,6 @@ start-time: null
 # End time as `YYYY-MM-DD`. Excludes datasets created after this time
 end-time: null
 data-limit: 0
-output-format: "messages" # "messages" or "supervised"
+output-format: "concat_dialog" # "messages" or "supervised" or "concat_dialog"
 savedir: "data/"
 # debug: True

src/download_and_filter.py

@@ -298,7 +298,7 @@ def main(args):
     parser.add(
         "-of", "--output-format", required=False,
         default="messages", type=str,
-        choices=["messages", "supervised"],
+        choices=["messages", "supervised", "concat_dialog"],
         help="The output format to save the data. By default it mimcs the format described in `preparers.py`. `supervised` means it saves as input-target pairs.")
     # Specify savedir
     parser.add(

src/downloader.py

@@ -1,7 +1,7 @@
 import os
 # import pandas as pd
 # from functools import partial
-from collections import defaultdict  # Counter,
+from collections import defaultdict, Counter
 # from datasets import load_dataset, list_datasets
 from helpers import io
 import random
@@ -106,6 +106,8 @@ def run_and_save(
         # If specified, reformat dataset for supervised learning, multi-turn dialogs, or reward modeling.
         if reformat == "supervised":
             prepared_dset = self._reformat_supervised(prepared_dset)
+        elif reformat == "concat_dialog":
+            prepared_dset = self._reformat_concat_dialog(prepared_dset)
 
         # save.
         savepath = os.path.join(savedir, f"{self.name}.jsonl.gz")
@@ -156,3 +158,28 @@ def dfs(node_id, parent_msg):
             dfs(root_id, dialog[root_id])
 
         return pairs
+
+    def _reformat_concat_dialog(self, dialogs):
+        reformatted = []
+        lens = []
+        for dialog in dialogs:
+            full_dialog = []
+            parent = None
+            for idx, message in enumerate(dialog):
+                if parent is None:
+                    full_dialog.append(message["text"])
+                    parent = 0
+                elif message["parent"] == parent:
+                    full_dialog.append(message["text"])
+                    parent = idx
+            lens.append(len(full_dialog))
+            reformatted.append({
+                "inputs": "\n".join(full_dialog),
+                "targets": "",
+                "dataset": dialog[0]["parent"],
+            })
+
+        # print(Counter(lens))
+        # print(reformatted[0])
+
+        return reformatted