From d903a5f48ca18eaf5cc65ff520282621bff428be Mon Sep 17 00:00:00 2001
From: Shayne Longpre <shayne@Shaynes-MacBook-Pro.local>
Date: Fri, 7 Mar 2025 14:22:45 -0500
Subject: [PATCH 1/2] fixed dialog

---
 src/download_and_filter.py |  2 +-
 src/downloader.py          | 29 ++++++++++++++++++++++++++++-
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/src/download_and_filter.py b/src/download_and_filter.py
index dfedee47..054e3e6b 100644
--- a/src/download_and_filter.py
+++ b/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(
diff --git a/src/downloader.py b/src/downloader.py
index 3ca51bf5..5567faee 100644
--- a/src/downloader.py
+++ b/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

From e62bc2cc6a21258c6fbfd6bd7f647d99cc5b16ad Mon Sep 17 00:00:00 2001
From: Shayne Longpre <shayne@Shaynes-MacBook-Pro.local>
Date: Fri, 7 Mar 2025 14:23:19 -0500
Subject: [PATCH 2/2] fixed

---
 src/configs/common_pile_ultra_permissive.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/configs/common_pile_ultra_permissive.yaml b/src/configs/common_pile_ultra_permissive.yaml
index d2de90a1..515761a4 100644
--- a/src/configs/common_pile_ultra_permissive.yaml
+++ b/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