Add vicuna dataset
diff --git a/model/model_training/custom_datasets/__init__.py b/model/model_training/custom_datasets/__init__.py
index 2b7e085442..f2620d6f8b 100644
--- a/model/model_training/custom_datasets/__init__.py
+++ b/model/model_training/custom_datasets/__init__.py
@@ -16,6 +16,7 @@
     QADataset,
     SODADialogue,
     TranslatedQA,
+    Vicuna,
     WebGPT,
 )
 from model_training.custom_datasets.rank_datasets import AugmentedOA
@@ -51,6 +52,7 @@
     "alpaca",
     "hf_summary",
     "hf_summary_pairs",
+    "vicuna",
 ]
 
 RM_DATASETS = [
@@ -137,6 +139,8 @@ def get_one_dataset(
     elif dataset_name == "oa_translated":
         # TODO make val_split lower..? by saganos
         dataset = TranslatedQA(data_path)
+    elif dataset_name == "vicuna":
+        dataset = Vicuna(cache_dir=data_path, **kwargs)
     elif dataset_name == "oasst_export":
         train, eval = load_oasst_export(data_path=data_path, val_split=val_split, mode=mode, **kwargs)
     elif dataset_name == "hf_summary":
diff --git a/model/model_training/custom_datasets/qa_datasets.py b/model/model_training/custom_datasets/qa_datasets.py
index d9eafccd43..e850a0b029 100644
--- a/model/model_training/custom_datasets/qa_datasets.py
+++ b/model/model_training/custom_datasets/qa_datasets.py
@@ -7,6 +7,7 @@
 import os
 import re
 from collections import defaultdict
+from pathlib import Path
 from urllib.request import urlopen
 
 import numpy as np
@@ -453,3 +454,58 @@ def __init__(self, mode: str = "sft", cache_dir: str = None) -> None:
 class CodeAlpaca(AlpacaBase):
     def __init__(self, mode: str = "sft", cache_dir: str = None) -> None:
         super().__init__(dataset_name="sahil2801/CodeAlpaca-20k", mode=mode, cache_dir=cache_dir)
+
+
+class Vicuna(Dataset):
+    name = "vicuna"
+
+    @staticmethod
+    def process_vicuna_conversations(data: list[dict[str, None | str]], input_max_length: int) -> list[str] | None:
+        dialogue = []
+        role = None
+        messages = []
+        # drop conversations that start with Bot
+        if len(data["conversations"]) == 0 or data["conversations"][0]["from"] != "human":
+            return None
+        for line in data["conversations"]:
+            speaker = line["from"]  # 'human' or 'gpt'
+            message = line["value"]
+            if role != speaker:
+                if role is not None:
+                    dialogue.append("\n".join(messages))
+                    messages = []
+                role = speaker
+            messages.append(message.strip())
+
+        if role is not None and len(messages) > 0:
+            dialogue.append("\n".join(messages))
+        dialogue_truncated = [k[:input_max_length] for k in dialogue]
+        return dialogue_truncated
+
+    def __init__(self, cache_dir: str | Path, mode: str = "sft", input_max_length: int = 2048) -> None:
+        super().__init__()
+
+        self.pairs = []
+        if mode not in ("sft", "rl"):
+            raise NotImplementedError(f"Currently only the modes 'sft' and 'rl' are implemented. Received {mode}.")
+        self.mode = mode
+        dataset = load_dataset(
+            "anon8231489123/ShareGPT_Vicuna_unfiltered",
+            cache_dir=cache_dir,
+            data_files=["ShareGPT_unfiltered_cleaned_split.json"],
+        )["train"]
+        for data in dataset:
+            if (
+                processed_data := self.process_vicuna_conversations(data, input_max_length=input_max_length)
+            ) is not None:
+                self.pairs.append(processed_data)
+
+    def __len__(self) -> int:
+        return len(self.pairs)
+
+    def __getitem__(self, index: int) -> list[str] | tuple[str]:
+        dialogue: list[str] = self.pairs[index]
+        if self.mode == "sft":
+            return dialogue
+        elif self.mode == "rl":
+            return tuple(dialogue[:-1])

#2341

This PR adds the vicuna dataset. We use the filtered vicuna dataset since it only contains english conversations.

The sanity checks are necessary here and filter quite a number of IMO unuseful conversations (where the conversation is empty, where we have less than 2 speakers and where the speakers are not in adjacent order). After all filterings we are left with 84260 of 85071 conversations.