Prepare oasst data export to HuggingFace (LAION-AI#2305)

Add conversion from trees to 'flat' message table jsonl file.

- add oasst-data function to read messages & message tree files
`read_message_trees`, `read_message_tree_list`, `read_messages`,
`read_message_list` (+ lower level functions `open_jsonl_read`,
`read_oasst_obj`, `read_oasst_jsonl`)
- add oasst-data fucntios to write messages & message tree files
`write_message_trees`, `write_messages` (+ lower level functions
`open_jsonl_write`, `write_tree`, `write_message`)
- add script used for data cleaning (`clean_dataset.py` in examples)
- add examples: `filter_trees.py`, `filter_messages.py`,
`split_dataset.py`, `tree_to_messages.py`
- update pre-commit hook black-jupyter to rev: 23.3.0 and add args:
["--profile", "black", "--filter-files"] to isort config

Author: andreaskoepf

.pre-commit-config.yaml

@@ -58,7 +58,7 @@ repos:
       - id: end-of-file-fixer
 
   - repo: https://github.com/psf/black
-    rev: 23.1.0
+    rev: 23.3.0
     hooks:
       - id: black-jupyter
 
@@ -71,6 +71,7 @@ repos:
     rev: 5.12.0
     hooks:
       - id: isort
+        args: ["--profile", "black", "--filter-files"]
 
   - repo: https://github.com/pre-commit/mirrors-prettier
     rev: v2.7.1

backend/oasst_backend/utils/tree_export.py

@@ -92,6 +92,7 @@ def prepare_export_message_node(
         message_id=message_id,
         parent_id=parent_id,
         user_id=user_id,
+        created_date=message.created_date,
         text=str(message.payload.payload.text),
         role=message.role,
         lang=message.lang,

model/model_training/custom_datasets/oasst_dataset.py

@@ -1,7 +1,7 @@
 from pathlib import Path
 from typing import Literal, Optional
 
-from oasst_data import ExportMessageNode, load_trees, visit_threads_depth_first
+from oasst_data import ExportMessageNode, read_message_trees, visit_threads_depth_first
 from torch import Generator
 from torch.utils.data import Dataset, random_split
 
@@ -43,7 +43,7 @@ def load_oasst_export(
         input_file_path = data_path / input_file_path
 
     threads_per_tree = []
-    for tree in load_trees(input_file_path):
+    for tree in read_message_trees(input_file_path):
         if tree.tree_state != "ready_for_export" or not tree.prompt.review_result or tree.prompt.lang not in lang_codes:
             continue
 

model/model_training/tools/check_oasst_export.py

@@ -1,6 +1,6 @@
 import argparse
 
-from oasst_data import ExportMessageTree, load_tree_list, visit_messages_depth_first
+from oasst_data import ExportMessageTree, read_message_tree_list, visit_messages_depth_first
 
 
 def parse_args():
@@ -31,7 +31,7 @@ def tree_filter(tree: ExportMessageTree) -> bool:
             and (lang_codes is None or tree.prompt.lang in lang_codes)
         )
 
-    trees = load_tree_list(args.input_file_path, filter=tree_filter)
+    trees = read_message_tree_list(args.input_file_path, filter=tree_filter)
     print(f"{len(trees)} trees")
 
     all_messages = []

oasst-data/examples/clean_dataset.py

@@ -0,0 +1,106 @@
+import argparse
+from collections import OrderedDict
+
+import pandas
+from oasst_data.reader import read_message_trees
+from oasst_data.schemas import ExportMessageNode, ExportMessageTree
+from oasst_data.traversal import visit_messages_depth_first
+from oasst_data.writer import write_message_trees
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="filter_dataset")
+    parser.add_argument(
+        "input_file_name",
+        type=str,
+        help="path to input .jsonl or .jsonl.gz input file",
+    )
+    parser.add_argument(
+        "output_file_name",
+        type=str,
+        help="path to output .jsonl or .jsonl.gz file",
+    )
+    parser.add_argument("--instructions", type=str, help="xlsx file with instructions")
+    parser.add_argument("--exclude-nulls", action="store_true", default=False)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+
+    instructions_df = pandas.read_excel(args.instructions, na_filter=False)
+
+    # load dataset and index messages by id
+    tree_by_id: dict[str, ExportMessageTree] = OrderedDict()
+    message_by_id: dict[str, ExportMessageNode] = {}
+
+    print(f"Reading: {args.input_file_name}")
+    for message_tree in read_message_trees(args.input_file_name):
+        tree_by_id[message_tree.message_tree_id] = message_tree
+
+        def index_message(msg: ExportMessageNode):
+            message_by_id[msg.message_id] = msg
+
+        visit_messages_depth_first(message_tree.prompt, index_message)
+
+    print(f"Loaded {len(tree_by_id)} trees with {len(message_by_id)} messages.")
+
+    def count_descendants(msg: ExportMessageNode):
+        i = 1
+        if msg.replies:
+            for r in msg.replies:
+                i += count_descendants(r)
+        return i
+
+    def delete_message(msg: ExportMessageNode):
+        if msg.parent_id is None:
+            tree_by_id.pop(msg.message_id)
+            print(f"Tree deleted: {msg.message_id}")
+        else:
+            parent_msg = message_by_id[msg.parent_id]
+            parent_msg.replies.remove(msg)
+            print(f"Branch deleted: {msg.message_id} ({count_descendants(msg)} messages)")
+
+    # cleaning
+    print("Cleaning...")
+    for index, row in instructions_df.iterrows():
+        id = row["UUID"]
+        msg = message_by_id.get(id)
+        if msg is None:
+            print(f"Not found: {id}")
+
+        action = row["Action"]
+        if action == "Delete":
+            print(f"deleting: {id}")
+            delete_message(msg)
+        elif action == "Replace":
+            print(f"replace: {id}")
+            replace = row["Replace"]
+            msg.text = replace
+        elif action == "Edit":
+            print(f"edit: {id}")
+            if row["Category"] == "Copy Code":
+                find = "\nCopy code\n"
+                replace = "\n\n"
+            else:
+                find = row["Find"]
+                replace = row["Replace"]
+            msg.text.index(find)  # make sure text is present
+            msg.text = msg.text.replace(find, replace)
+        else:
+            print(f"Unsupported action {action}")
+
+    print("Done")
+
+    # write cleaned dataset to output file
+    print(f"Writing: {args.output_file_name}")
+    write_message_trees(
+        args.output_file_name,
+        tree_by_id.values(),
+        exclude_none=args.exclude_nulls,
+    )
+
+
+if __name__ == "__main__":
+    main()

oasst-data/examples/filter_messages.py

@@ -0,0 +1,153 @@
+import argparse
+import json
+
+from oasst_data import read_message_list, write_messages
+from oasst_data.schemas import ExportMessageNode
+from oasst_data.writer import open_jsonl_write
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="filter_messages")
+    parser.add_argument(
+        "input_file_name",
+        type=str,
+        help="path to input .jsonl or .jsonl.gz input file",
+    )
+    parser.add_argument(
+        "output_file_name",
+        type=str,
+        help="path to output .jsonl or .jsonl.gz file",
+    )
+    parser.add_argument(
+        "--include-deleted",
+        action="store_true",
+        help="Include deleted messages in export",
+    )
+    parser.add_argument(
+        "--deleted-only",
+        action="store_true",
+        help="Export only deleted messages (implies --include-deleted)",
+    )
+    parser.add_argument(
+        "--include-spam",
+        action="store_true",
+        help="Export including messages with no review or negative review result.",
+    )
+    parser.add_argument(
+        "--spam-only",
+        action="store_true",
+        help="Export only messages with negative review result (implies --include-spam).",
+    )
+    parser.add_argument(
+        "--exclude-normal",
+        action="store_true",
+        help="exclude non-deleted non-synthetic messages with positive review",
+        default=False,
+    )
+    parser.add_argument(
+        "--include-synthetic",
+        action="store_true",
+        help="Include synthetic messages in export",
+    )
+    parser.add_argument(
+        "--synthetic-only",
+        action="store_true",
+        help="Export only synthetic messages (implies --include-synth)",
+    )
+    parser.add_argument(
+        "--user",
+        type=str,
+        help="Only export trees involving the user with the specified ID. Incompatible with --state.",
+    )
+    parser.add_argument(
+        "--state",
+        type=str,
+        help="all|prompt_lottery_waiting|growing|ready_for_export|aborted_low_grade|halted_by_moderator|backlog_ranking",
+    )
+    parser.add_argument(
+        "--lang",
+        type=str,
+        help="Filter message trees by language code (BCP 47)",
+    )
+    parser.add_argument(
+        "--prompts-only",
+        action="store_true",
+        help="Export a list of initial prompt messages",
+    )
+    parser.add_argument(
+        "--export-text-only",
+        action="store_true",
+        help="Write jsonl file with message text strings only",
+    )
+    parser.add_argument("--exclude-nulls", action="store_true", default=False)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+
+    deleted: bool | None = False
+    spam: bool | None = False
+    synthetic: bool | None = False
+    langs: list[str] | None = None
+    states: list[str] | None = None
+    prompts_only: bool = args.prompts_only
+    exclude_normal: bool = args.exclude_normal
+
+    if args.include_deleted:
+        deleted = None
+    elif args.deleted_only:
+        deleted = True
+
+    if args.include_spam:
+        spam = None
+    elif args.spam_only:
+        spam = True
+
+    if args.include_synthetic:
+        synthetic = None
+    elif args.synthetic_only:
+        synthetic = True
+
+    if args.lang:
+        langs = args.lang.split(",")
+
+    if args.state:
+        states = args.state.split(",")
+
+    def approve_message(msg: ExportMessageNode) -> bool:
+        if (
+            (deleted is not None and msg.deleted != deleted)
+            or (synthetic is not None and msg.synthetic != synthetic)
+            or (prompts_only and msg.parent_id)
+            or (langs is not None and msg.lang not in langs)
+            or (states is not None and msg.tree_state not in states)
+        ):
+            return False
+
+        if exclude_normal is True and not msg.deleted and not msg.synthetic and msg.review_result:
+            return False
+
+        if spam is not None and spam != (not msg.review_result):
+            return False
+
+        return True
+
+    print(f"Reading: {args.input_file_name}")
+    messages = read_message_list(args.input_file_name, approve_message)
+
+    print(f"Found {len(messages)} matching messages.")
+
+    print(f"Writing: {args.output_file_name}")
+    if args.export_text_only:
+        with open_jsonl_write(args.output_file_name) as file:
+            for msg in messages:
+                json.dump(msg.text, file)
+                file.write("\n")
+    else:
+        write_messages(args.output_file_name, messages, args.exclude_nulls)
+
+
+if __name__ == "__main__":
+    main()

oasst-data/examples/filter_trees.py

@@ -0,0 +1,56 @@
+import argparse
+
+from oasst_data import read_message_trees, write_message_trees
+from oasst_data.schemas import ExportMessageTree
+from oasst_data.traversal import visit_messages_depth_first
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="filter_tres")
+    parser.add_argument(
+        "input_file_name",
+        type=str,
+        help="path to input .jsonl or .jsonl.gz input file",
+    )
+    parser.add_argument(
+        "output_file_name",
+        type=str,
+        help="path to output .jsonl or .jsonl.gz file",
+    )
+    parser.add_argument(
+        "--states",
+        type=str,
+        default="ready_for_export",
+        help="all|prompt_lottery_waiting|growing|ready_for_export|aborted_low_grade|halted_by_moderator|backlog_ranking",
+    )
+    parser.add_argument("--exclude-nulls", action="store_true", default=False)
+    parser.add_argument("--allow-synth", action="store_true", default=False)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+
+    # load dataset and index messages by id
+    trees: list[ExportMessageTree] = []
+
+    states = args.states.split(",")
+    allow_synth = args.allow_synth
+
+    print(f"Reading: {args.input_file_name}")
+    for message_tree in read_message_trees(args.input_file_name):
+        msgs = []
+        visit_messages_depth_first(message_tree.prompt, msgs.append)
+        if message_tree.tree_state in states:
+            if allow_synth or not any(x.synthetic for x in msgs):
+                trees.append(message_tree)
+
+    print(f"Found {len(trees)} matching trees.")
+
+    print(f"Writing: {args.output_file_name}")
+    write_message_trees(args.output_file_name, trees, exclude_none=args.exclude_nulls)
+
+
+if __name__ == "__main__":
+    main()