New oasst export dataset loader
diff --git a/backend/import.py b/backend/import.py
index 2880c34412..ecea434bab 100644
--- a/backend/import.py
+++ b/backend/import.py
@@ -16,7 +16,7 @@
 from oasst_backend.models.payload_column_type import PayloadContainer
 from oasst_backend.prompt_repository import PromptRepository
 from oasst_backend.user_repository import UserRepository
-from oasst_backend.utils.tree_export import ExportMessageNode, ExportMessageTree
+from oasst_shared.schemas.export import ExportMessageNode, ExportMessageTree
 from sqlmodel import Session
 
 # well known id
diff --git a/backend/oasst_backend/utils/tree_export.py b/backend/oasst_backend/utils/tree_export.py
index 73beed9b29..06f7a0d55b 100644
--- a/backend/oasst_backend/utils/tree_export.py
+++ b/backend/oasst_backend/utils/tree_export.py
@@ -11,54 +11,25 @@
 from fastapi.encoders import jsonable_encoder
 from oasst_backend.models import Message
 from oasst_backend.models.message_tree_state import State as TreeState
-from pydantic import BaseModel
-
-
-class LabelAvgValue(BaseModel):
-    value: float | None
-    count: int
-
-
-LabelValues = dict[str, LabelAvgValue]
-
-
-class ExportMessageNode(BaseModel):
-    message_id: str
-    parent_id: str | None
-    text: str
-    role: str
-    lang: str | None
-    review_count: int | None
-    review_result: bool | None
-    rank: int | None
-    synthetic: bool | None
-    model_name: str | None
-    emojis: dict[str, int] | None
-    replies: list[ExportMessageNode] | None
-    labels: LabelValues | None
-
-    @staticmethod
-    def prep_message_export(message: Message, labels: Optional[LabelValues] = None) -> ExportMessageNode:
-        return ExportMessageNode(
-            message_id=str(message.id),
-            parent_id=str(message.parent_id) if message.parent_id else None,
-            text=str(message.payload.payload.text),
-            role=message.role,
-            lang=message.lang,
-            review_count=message.review_count,
-            review_result=message.review_result if message.review_result or message.review_count > 2 else None,
-            synthetic=message.synthetic,
-            model_name=message.model_name,
-            emojis=message.emojis,
-            rank=message.rank,
-            labels=labels,
-        )
-
-
-class ExportMessageTree(BaseModel):
-    message_tree_id: str
-    tree_state: Optional[str]
-    prompt: Optional[ExportMessageNode]
+from oasst_shared.schemas.export import ExportMessageNode, ExportMessageTree, LabelValues
+
+
+def prepare_export_message_node(message: Message, labels: Optional[LabelValues] = None) -> ExportMessageNode:
+    return ExportMessageNode(
+        message_id=str(message.id),
+        parent_id=str(message.parent_id) if message.parent_id else None,
+        text=str(message.payload.payload.text),
+        role=message.role,
+        lang=message.lang,
+        deleted=message.deleted,
+        review_count=message.review_count,
+        review_result=message.review_result if message.review_result or message.review_count > 2 else None,
+        synthetic=message.synthetic,
+        model_name=message.model_name,
+        emojis=message.emojis,
+        rank=message.rank,
+        labels=labels,
+    )
 
 
 def build_export_tree(
@@ -67,7 +38,7 @@ def build_export_tree(
     messages: list[Message],
     labels: Optional[dict[UUID, LabelValues]] = None,
 ) -> ExportMessageTree:
-    export_messages = [ExportMessageNode.prep_message_export(m, labels.get(m.id) if labels else None) for m in messages]
+    export_messages = [prepare_export_message_node(m, labels.get(m.id) if labels else None) for m in messages]
 
     messages_by_parent = defaultdict(list)
     for message in export_messages:
@@ -133,7 +104,7 @@ def write_messages_to_file(
 
     with out_buff as f:
         for m in messages:
-            export_message = ExportMessageNode.prep_message_export(m, labels.get(m.id) if labels else None)
+            export_message = prepare_export_message_node(m, labels.get(m.id) if labels else None)
 
             file_data = jsonable_encoder(export_message, exclude_none=True)
             json.dump(file_data, f)
diff --git a/model/model_training/.gitignore b/model/model_training/.gitignore
index 5c5c0c1ccc..2b395ff822 100644
--- a/model/model_training/.gitignore
+++ b/model/model_training/.gitignore
@@ -1,2 +1,3 @@
 .cache
 wandb
+saved_model
diff --git a/model/model_training/configs/config.yaml b/model/model_training/configs/config.yaml
index dcb12e0e9e..5f4bf64f09 100644
--- a/model/model_training/configs/config.yaml
+++ b/model/model_training/configs/config.yaml
@@ -54,11 +54,12 @@ defaults:
 
 oa_dataset_only:
   datasets:
-    - oa_private:
-        split: sft
-        val_split: 0.0
-        fraction: 1
-        file: 2023-02-10_oasst_prod.jsonl
+    - oasst_export:
+        lang: "en,es,de,fr"
+        top_k: 2
+        input_file_path: 2023-02-19_oasst_ready_with_spam_deleted.jsonl.gz
+  num_train_epochs: 10
+  save_steps: 2000
 
 pythia:
   learning_rate: 8e-6
@@ -72,6 +73,17 @@ pythia:
   per_device_eval_batch_size: 4
   output_dir: pythia_model
 
+pythia-1B:
+  learning_rate: 8e-6
+  model_name: EleutherAI/pythia-1b-deduped
+  weight_decay: 0.01
+  max_length: 520
+  warmup_steps: 1000
+  gradient_checkpointing: false
+  gradient_accumulation_steps: 2
+  per_device_train_batch_size: 16
+  per_device_eval_batch_size: 32
+
 galactica-125m:
   learning_rate: 5e-5
   model_name: facebook/galactica-125m
diff --git a/model/model_training/custom_datasets/__init__.py b/model/model_training/custom_datasets/__init__.py
index b3d2cc7fa1..b8bbdcad37 100644
--- a/model/model_training/custom_datasets/__init__.py
+++ b/model/model_training/custom_datasets/__init__.py
@@ -1,6 +1,7 @@
 """
     High level functions for model training
 """
+from custom_datasets.oasst_dataset import load_oasst_export
 from custom_datasets.prompt_dialogue import OAPrivate, PrivateInstructionTuning
 from custom_datasets.qa_datasets import SODA, JokeExplaination, QADataset, SODADialogue, TranslatedQA, WebGPT
 from custom_datasets.summarization import SummarizationDataset
@@ -82,6 +83,8 @@ def get_one_dataset(conf, dataset_name, val_split=0.2, data_path=None, mode="sft
         dataset = TranslatedQA(data_path)
     elif dataset_name == "oa_private":
         dataset = OAPrivate(data_path, **kwargs)
+    elif dataset_name == "oasst_export":
+        train, eval = load_oasst_export(data_path=data_path, val_split=val_split, **kwargs)
     else:
         raise ValueError(f"Unknown dataset {dataset_name}")
 
diff --git a/model/model_training/custom_datasets/oasst_dataset.py b/model/model_training/custom_datasets/oasst_dataset.py
new file mode 100644
index 0000000000..ada2413331
--- /dev/null
+++ b/model/model_training/custom_datasets/oasst_dataset.py
@@ -0,0 +1,124 @@
+import gzip
+import json
+from pathlib import Path
+from typing import Callable, Optional
+
+import pydantic
+from custom_datasets.formatting import format_pair
+from oasst_shared.schemas.export import ExportMessageNode, ExportMessageTree
+from torch import Generator, default_generator
+from torch.utils.data import Dataset, random_split
+
+
+def _visit_threads_depth_first(
+    node: ExportMessageNode,
+    visitor: Callable[[list[ExportMessageNode]], None],
+    predicate: Optional[Callable[[list[ExportMessageNode]], bool]] = None,
+    parents: list[ExportMessageNode] = None,
+):
+    parents = parents or []
+    if not node:
+        return
+    thread = parents + [node]
+    if predicate is None or predicate(thread):
+        visitor(thread)
+    if node.replies:
+        parents = thread
+        for c in node.replies:
+            _visit_threads_depth_first(node=c, visitor=visitor, predicate=predicate, parents=parents)
+
+
+class ListDataset(Dataset):
+    def __init__(self, data: list):
+        super().__init__()
+        self.data = data
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        return self.data[index]
+
+
+def load_oasst_export(
+    input_file_path: str | Path,
+    val_split: float = 0.2,
+    lang: str = "en",
+    top_k: Optional[int] = None,
+    generator: Optional[Generator] = default_generator,
+    data_path: str | Path = None,
+) -> tuple[ListDataset, ListDataset]:
+    lang_codes = lang.split(",")
+
+    if not isinstance(input_file_path, Path):
+        input_file_path = Path(input_file_path)
+    if not input_file_path.is_absolute() and data_path:
+        if not isinstance(data_path, Path):
+            data_path = Path(data_path)
+        input_file_path = data_path / input_file_path
+
+    if input_file_path.suffix == ".gz":
+        file_in = gzip.open(str(input_file_path), mode="tr", encoding="UTF-8")
+    else:
+        file_in = input_file_path.open("r", encoding="UTF-8")
+
+    threads_per_tree = []
+
+    with file_in:
+        # read one message tree per line
+        for line in file_in:
+            dict_tree = json.loads(line)
+
+            # validate data
+            tree: ExportMessageTree = pydantic.parse_obj_as(ExportMessageTree, dict_tree)
+
+            if (
+                tree.tree_state != "ready_for_export"
+                or not tree.prompt.review_result
+                or tree.prompt.lang not in lang_codes
+            ):
+                continue
+
+            # extract all threads up to last asssitant reply
+            threads: list[list[ExportMessageNode]] = []
+
+            def thread_filter(thread: list[ExportMessageNode]) -> bool:
+                if any(m.deleted for m in thread):
+                    return False
+
+                if top_k is not None:
+                    for i, m in enumerate(thread):
+                        if m.role == "assistant":
+                            if m.rank is None:
+                                if len(thread[i - 1].replies) > 1:
+                                    return False
+                            elif m.rank >= top_k:
+                                return False
+                return True
+
+            def leaf_filter(thread: list[ExportMessageNode]) -> bool:
+                return (
+                    len(thread) > 1
+                    and not thread[-1].replies
+                    and (thread[-1].role == "assistant" or thread[-2].replies[0] == thread[-1])
+                    and thread_filter(thread)
+                )
+
+            _visit_threads_depth_first(tree.prompt, threads.append, leaf_filter)
+            for t in threads:
+                if t[-1].role == "prompter":
+                    t.pop()
+
+            threads_per_tree.append(threads)
+
+    # split on tree basis, messages from same tree must not end up in different splits
+    trees = ListDataset(threads_per_tree)
+    splits = random_split(trees, lengths=[1.0 - val_split, val_split], generator=generator)
+
+    def flatten(d: ListDataset) -> ListDataset:
+        return ListDataset([format_pair([m.text for m in t]) for ts in d for t in ts])
+
+    train = flatten(splits[0])
+    val = flatten(splits[1])
+
+    return train, val
diff --git a/model/model_training/tests/test_oasst_dataset.py b/model/model_training/tests/test_oasst_dataset.py
new file mode 100644
index 0000000000..5101d14410
--- /dev/null
+++ b/model/model_training/tests/test_oasst_dataset.py
@@ -0,0 +1,18 @@
+from argparse import Namespace
+
+from custom_datasets import get_one_dataset
+
+
+def test_load_oasst_export_dataset():
+    config = Namespace(
+        cache_dir=".cache",
+    )
+    kwargs = {
+        "lang": "en,es,de,fr",
+        "top_k": 2,
+        "input_file_path": "2023-02-19_oasst_ready_with_spam_deleted.jsonl.gz",
+    }
+    print("taeiae")
+    train, val = get_one_dataset(conf=config, dataset_name="oasst_export", **kwargs)
+    assert len(train) > 9000
+    assert len(val) > 2000
diff --git a/oasst-shared/oasst_shared/schemas/export.py b/oasst-shared/oasst_shared/schemas/export.py
new file mode 100644
index 0000000000..ac2aa14403
--- /dev/null
+++ b/oasst-shared/oasst_shared/schemas/export.py
@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class LabelAvgValue(BaseModel):
+    value: float | None
+    count: int
+
+
+LabelValues = dict[str, LabelAvgValue]
+
+
+class ExportMessageNode(BaseModel):
+    message_id: str
+    parent_id: str | None
+    text: str
+    role: str
+    lang: str | None
+    review_count: int | None
+    review_result: bool | None
+    deleted: bool | None
+    rank: int | None
+    synthetic: bool | None
+    model_name: str | None
+    emojis: dict[str, int] | None
+    replies: list[ExportMessageNode] | None
+    labels: LabelValues | None
+
+
+class ExportMessageTree(BaseModel):
+    message_tree_id: str
+    tree_state: Optional[str]
+    prompt: Optional[ExportMessageNode]

- support jsonl and jsonl.gz input files
- use oasst export classes for parsing (classes reside now in oasst_shared)
- extract all usable leaf-nodes (last assistant replies of conversation threads)
- allow filtering by language and top_k (ranking results)
- split into train/eval (while ensuring that no conversation threads of the same tree end up both in train & eval)