add dataset counts script
diff --git a/model/model_training/README.md b/model/model_training/README.md
index d9f02c44fa..da5a338ac3 100644
--- a/model/model_training/README.md
+++ b/model/model_training/README.md
@@ -164,6 +164,41 @@ the end to trigger deepspeed
 python trainer_sft.py --configs defaults your-model-name --deepspeed
 ```
 
+### Datasets
+
+Here is an uncomplete overview of datasets for sft:
+
+<!-- prettier-ignore -->
+dataset_name        | train_counts | eval_counts | total_counts
+----------------------------------------------------------------
+
+<!-- prettier-ignore -->
+webgpt              |     15662    |     3916    |     19578
+squad_v2            |    130319    |    11873    |    142192
+adversarial_qa      |     30000    |     3000    |     33000
+trivia_qa_nocontext |    138384    |    17944    |    156328
+xsum                |    204045    |    11332    |    215377
+cnn_dailymail       |    287113    |    13368    |    300481
+multi_news          |     44972    |     5622    |     50594
+scitldr             |      1992    |      619    |      2611
+joke                |       301    |       76    |       377
+gsm8k               |      7473    |     1319    |      8792
+dive_mt             |      6192    |     1548    |      7740
+
+This list can be generated with the following command, but beware that this
+downloads all available datasets (>100GB):
+
+```bash
+python check_dataset_counts.py --datasets all --mode sft
+```
+
+One can specify datasets, which can be found in the config corresponding to the
+mode the mode (e.g. configs/config.yaml for sft, configs/config_rm.yaml for rm):
+
+```bash
+python check_dataset_counts.py --datasets webgpt squad_v2 --mode sft
+```
+
 ### Troubleshooting
 
 - If training on a VM, you might need to install OpenMPI. Check out
diff --git a/model/model_training/check_dataset_counts.py b/model/model_training/check_dataset_counts.py
new file mode 100644
index 0000000000..e4853f4256
--- /dev/null
+++ b/model/model_training/check_dataset_counts.py
@@ -0,0 +1,125 @@
+import argparse
+from enum import Enum
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+import yaml
+from model_training.utils import _strtobool, get_dataset
+
+
+class Mode(Enum):
+    sft = "sft"
+    rm = "rm"
+    rl = "rl"
+
+    def config_name(self) -> str:
+        match self:
+            case Mode.sft:
+                return "config.yaml"
+            case Mode.rm:
+                return "config_rm.yaml"
+            case Mode.rl:
+                return "config_rl.yaml"
+
+    def default_config(self) -> str:
+        match self:
+            case Mode.sft:
+                return "defaults"
+            case Mode.rm:
+                return "defaults_rm"
+            case Mode.rl:
+                return "defaults_rlhf"
+
+
+def read_yaml(dir: str | Path, config_file: str) -> dict[str, Any]:
+    with open(Path(dir) / config_file, "r") as f:
+        return yaml.safe_load(f)
+
+
+def argument_parsing(notebook=False, notebook_args=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--datasets",
+        nargs="+",
+        required=True,
+        help="""
+        Multiple datasets can be passed to set different options.
+        For example, run as:
+
+           ./check_dataset_counts.py --datasets math oasst_export_eu
+
+        to check the counts of the math and the oasst_export_eu dataset.
+    """,
+    )
+    parser.add_argument("--mode", dest="mode", type=Mode, choices=list(Mode))
+    parser.add_argument("--output_path", dest="output_path", default="dataset_counts.csv")
+
+    if notebook:
+        args, remaining = parser.parse_known_args(notebook_args)
+    else:
+        args, remaining = parser.parse_known_args()
+
+    # Config from YAML
+    mode: Mode = args.mode
+    configs = read_yaml("./configs", config_file=mode.config_name())
+    conf = configs[mode.default_config()]
+    if "all" in args.datasets:
+        conf["datasets"] = configs[mode.default_config()]["datasets"] + configs[mode.default_config()]["datasets_extra"]
+    else:
+        # reset datasets, so that we only get the datasets defined in configs and remove the ones in the default
+        datasets_list = list()
+
+        for name in args.datasets:
+            # check and process multiple datasets
+            if "," in name:
+                for n in name.split(","):
+                    datasets_value = configs[n].get("datasets") or configs[n]["datasets_extra"]
+            # check if dataset is extra key in config
+            elif name in configs:
+                datasets_value = configs[name].get("datasets") or configs[name]["datasets_extra"]
+            # check in default config
+            elif name in configs[mode.default_config()]["datasets"]:
+                datasets_value = [name]
+            else:
+                raise ValueError(
+                    f'Error: Could not find the dataset "{name}" in {mode.config_name()}. ',
+                    f"Tried to look for this dataset within th key {mode.default_config()} ",
+                    "and as seperate key.",
+                )
+
+            datasets_list.extend(datasets_value)
+    conf["mode"] = mode
+    conf["output_path"] = args.output_path
+    conf["datasets_extra"] = []
+    conf["datasets"] = datasets_list
+    # Override config from command-line
+    parser = argparse.ArgumentParser()
+    for key, value in conf.items():
+        type_ = type(value) if value is not None else str
+        if type_ == bool:
+            type_ = _strtobool
+        parser.add_argument(f"--{key}", type=type_, default=value)
+        # Allow --no-{key}  to remove it completely
+        parser.add_argument(f"--no-{key}", dest=key, action="store_const", const=None)
+
+    args = parser.parse_args(remaining)
+    print(args)
+    return args
+
+
+if __name__ == "__main__":
+    args = argument_parsing()
+
+    train, evals = get_dataset(args, mode=args.mode)
+    overview_df = pd.DataFrame(columns=["dataset_name", "train_counts", "eval_counts", "total_counts"])
+    for idx, dataset_name in enumerate(args.datasets):
+        eval_count = len(evals.get(dataset_name, []))
+        overview_df.loc[idx] = [
+            dataset_name,
+            len(train.datasets[idx]),
+            eval_count,
+            len(train.datasets[idx]) + eval_count,
+        ]
+    print(overview_df)
+    overview_df.to_csv(args.output_path, index=False)
diff --git a/model/model_training/configs/config.yaml b/model/model_training/configs/config.yaml
index bfa3456a67..d6c82a76ff 100644
--- a/model/model_training/configs/config.yaml
+++ b/model/model_training/configs/config.yaml
@@ -29,7 +29,7 @@ defaults:
     - cnn_dailymail
     - multi_news
     - scitldr
-    - soda
+    # - soda  # TODO: fix list index out of range error
     - joke
     - gsm8k
     - dive_mt
@@ -41,16 +41,16 @@ defaults:
     - wmt2019_de-en
     - samsum
     - soda_dialogue
-    # instructional_datasets:
-    # - humaneval_mbpp_codegen_qa
-    # - humaneval_mbpp_testgen_qa
-    # - grade_school_math_instructions
-    # - recipes
-    # - ubuntu_dialogue_qa
-    # - cmu_wiki_qa
-    # - youtube_subs_howto100M
-    # - iapp_wiki_qa_squad
-    # - zhihu-kol
+  # instructional_datasets:
+  #  - humaneval_mbpp_codegen_qa
+  #  - humaneval_mbpp_testgen_qa
+  #  - grade_school_math_instructions
+  #  - recipes
+  #  - ubuntu_dialogue_qa
+  #  - cmu_wiki_qa
+  #  - youtube_subs_howto100M
+  #  - iapp_wiki_qa_squad
+  #  - zhihu-kol
   datasets_extra: [] # For config options to add additional datasets, since yaml doesn't let us extend arrays
   cache_dir: .cache
   loss_fn: CrossEntropyLoss
diff --git a/model/model_training/utils.py b/model/model_training/utils.py
index a87bbda993..d813db6272 100644
--- a/model/model_training/utils.py
+++ b/model/model_training/utils.py
@@ -354,7 +354,10 @@ def get_dataset_name_and_kwargs_from_data_config(data_config):
         return data_config, {}
 
 
-def get_dataset(conf, mode="sft"):
+def get_dataset(
+    conf,
+    mode: str = "sft",
+) -> tuple[ConcatDataset, dict[str, Subset]]:
     train_datasets, evals = [], {}
 
     for data_config in conf.datasets + conf.datasets_extra:

This can be used for a part of #2174.

This PR contains:
 - a generic script to count the dataset entries for any dataset and any mode
 - updated model/model_training/README.md with some counts of datasets
 - some minor adjustments to model/model_training/utils.py so that the datasets can be returned as list. This is done so that this PR can reuse existing code and does not need to copy+paste.

