more data refactor

src/ctx_to_lora/data/collator.py

@@ -0,0 +1,58 @@
+import torch
+from transformers.data import DataCollatorWithFlattening
+
+flattener = DataCollatorWithFlattening()
+
+
+def train_packed_collator(inp_list):
+    # no padding
+    packed_inputs = flattener(inp_list, return_tensors="pt")
+    if "ctx_ids" in inp_list[0]:
+        ctx_ids = [{"input_ids": example["ctx_ids"]} for example in inp_list]
+        packed_ctx = flattener(ctx_ids, return_tensors="pt")
+        packed_inputs["ctx_ids"] = packed_ctx["input_ids"]
+        packed_inputs["ctx_position_ids"] = packed_ctx["position_ids"]
+    return packed_inputs
+
+
+def train_collator(inp_list, tokenizer):
+    # input is a list of tokenized sequences
+    padding_kwargs = dict(
+        padding=True,
+        padding_side="right",
+        pad_to_multiple_of=8,
+        return_tensors="pt",
+    )
+
+    ctx_ids = None
+    if "ctx_ids" in inp_list[0]:
+        # have to be manual since it has [ctx_len, features] shape
+        # pad to the longest ctx_len in the batch
+        # which can have a different length from the input_ids, attn_mask, labels
+
+        ctx_ids = [example.pop("ctx_ids") for example in inp_list]
+        ctx_ids = torch.nn.utils.rnn.pad_sequence(
+            ctx_ids,
+            batch_first=True,
+            padding_value=0,
+        )
+        # exotic keys won't be padded, so we need to pad them as well
+        ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
+        ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
+            ctx_attn_mask,
+            batch_first=True,
+            padding_value=0,
+        )
+
+    labels = [x.pop("labels") for x in inp_list]
+    padded_seq = tokenizer.pad(inp_list, **padding_kwargs)
+
+    # hacky explicit padding since the labels are not padded by default
+    labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"]
+    labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels)
+    out = {**padded_seq, "labels": labels}
+    if ctx_ids is not None:
+        out["ctx_ids"] = ctx_ids
+        out["ctx_attn_mask"] = ctx_attn_mask
+
+    return out

src/ctx_to_lora/data/definitions.py

@@ -0,0 +1,377 @@
+from glob import glob
+
+IGNORE_INDEX = -100
+
+FW_QA_PATHS = [
+    f"data/raw_datasets/fw_qa/{i:05d}.parquet" for i in [0, 1, 6, 7, 8, 10, 22, 30, 35]
+]
+
+TRANSFORMED_DATA_DIR = "data/processed_datasets"
+
+# approximate length of the datasets (train split)
+# needed for streaming datasets
+DS_LEN = {
+    "fw_qa_3_mini": 100_000,
+    "fw_qa_3_medium": 121_000_000,
+    "fw_qa_3": 270_000_000,
+    "fw_qa_xl": 27_000_000,
+    "ctx_qa": 300_000,
+    "pwc": 240_000,
+    "hotpot_qa": 100_000,
+    "squad": 90_000,
+    "drop": 77_000,
+    "narrativeqa": 40_000,
+    "quoref": 11_000,
+    "ropes": 11_000,
+    "synthetic_convqa": 40_000,
+}
+
+LONGBENCH_TASKS = [
+    "longbench/narrativeqa",
+    "longbench/qasper",
+    "longbench/multifieldqa_en",
+    "longbench/multifieldqa_zh",
+    "longbench/hotpotqa",
+    "longbench/2wikimqa",
+    "longbench/musique",
+    "longbench/dureader",
+    "longbench/gov_report",
+    "longbench/qmsum",
+    "longbench/multi_news",
+    "longbench/vcsum",
+]
+
+LONGBENCH_E_TASKS = [
+    "longbench/qasper_e",
+    "longbench/multifieldqa_en_e",
+    "longbench/hotpotqa_e",
+    "longbench/2wikimqa_e",
+    "longbench/gov_report_e",
+    "longbench/multi_news_e",
+]
+
+DS_KWARGS = {
+    "hotpot_qa": dict(
+        train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[900:]"),
+        validation=dict(
+            path="hotpotqa/hotpot_qa", name="fullwiki", split="train[:900]"
+        ),
+        test=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="validation"),
+    ),
+    "hotpot_qa_tiny": dict(
+        train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[900:2000]"),
+        validation=dict(
+            path="hotpotqa/hotpot_qa", name="fullwiki", split="train[:900]"
+        ),
+    ),
+    "pwc": dict(
+        train=dict(path="sggetao/PwC", split="train[900:]"),
+        validation=dict(path="sggetao/PwC", split="train[:900]"),
+        test=dict(path="sggetao/PwC", split="test"),
+    ),
+    "pwc_tiny": dict(
+        train=dict(path="sggetao/PwC", split="train[900:2000]"),
+        validation=dict(path="sggetao/PwC", split="train[:900]"),
+    ),
+    "squad": dict(
+        train=dict(path="rajpurkar/squad", split="train[900:]"),
+        validation=dict(path="rajpurkar/squad", split="train[:900]"),
+        test=dict(path="rajpurkar/squad", split="validation"),
+    ),
+    # "fw_qa_tiny": dict(
+    #     train=dict(
+    #         path="parquet",
+    #         data_files="data/raw_datasets/fw_qa/00000.parquet",
+    #         split="train",
+    #     ),
+    #     validation=dict(
+    #         path="parquet",
+    #         data_files="data/raw_datasets/fw_qa/00000_val.parquet",
+    #         split="train",
+    #     ),
+    # ),
+    # "fw_qa": dict(
+    #     train=dict(
+    #         path="parquet",
+    #         data_files=FW_QA_PATHS,
+    #         split="train",
+    #     ),
+    #     validation=dict(
+    #         path="parquet",
+    #         data_files="data/raw_datasets/fw_qa/*_val.parquet",
+    #         split="train",
+    #     ),
+    # ),
+    # "fw_qa_large": dict(
+    #     train=dict(
+    #         path="parquet",
+    #         data_files=glob("data/raw_datasets/fw_qa_large/*res.parquet"),
+    #         split="train",
+    #     ),
+    #     validation=dict(
+    #         path="parquet",
+    #         data_files=glob("data/raw_datasets/fw_qa_large/*val.parquet"),
+    #         split="train",
+    #     ),
+    # ),
+    "fw_qa_xl": dict(
+        train=dict(
+            path="parquet",
+            data_files=glob("data/raw_datasets/fw_qa_xl/*[!val].parquet"),
+            split="train",
+        ),
+        validation=dict(
+            path="parquet",
+            data_files=glob("data/raw_datasets/fw_qa_xl/*val.parquet"),
+            split="train",
+        ),
+        # there is no explcit test for this ds
+        # used only for sanity check
+        test=dict(
+            path="parquet",
+            data_files=glob("data/raw_datasets/fw_qa_xl/*val.parquet"),
+            split="train",
+        ),
+    ),
+    # "fw_qa_2": dict(
+    #     train=dict(
+    #         path="parquet",
+    #         data_files=glob("data/raw_datasets/fw_qa_2/*[!val].parquet"),
+    #         split="train",
+    #     ),
+    #     validation=dict(
+    #         path="parquet",
+    #         data_files=glob("data/raw_datasets/fw_qa_2/*val.parquet"),
+    #         split="train",
+    #     ),
+    # ),
+    "fw_qa_3_mini": dict(
+        train=dict(
+            path="parquet",
+            data_files=glob("data/raw_datasets/fw_qa_3/000_00001.parquet"),
+            split="train[:100000]",
+        ),
+    ),
+    "fw_qa_3_medium": dict(
+        train=dict(
+            path="parquet",
+            data_files=glob("data/raw_datasets/fw_qa_3/00[0-5]*[!val].parquet"),
+            split="train",
+        ),
+    ),
+    "fw_qa_3": dict(
+        train=dict(
+            path="parquet",
+            data_files=glob("data/raw_datasets/fw_qa_3/*[!val].parquet"),
+            split="train",
+        ),
+        validation=dict(
+            path="parquet",
+            data_files=glob("data/raw_datasets/fw_qa_3/*val.parquet"),
+            split="train",
+        ),
+    ),
+    "ctx_qa": dict(
+        train=dict(
+            path="parquet",
+            data_files=glob("data/raw_datasets/ctx_qa/*res.parquet"),
+            split="train",
+        ),
+        validation=dict(
+            path="parquet",
+            data_files=glob("data/raw_datasets/ctx_qa/*val.parquet"),
+            split="train",
+        ),
+        # there is no explcit test for this ds
+        # used only for sanity check
+        test=dict(
+            path="parquet",
+            data_files=glob("data/raw_datasets/ctx_qa/*val.parquet"),
+            split="train",
+        ),
+    ),
+    "drop": dict(
+        train=dict(
+            path="ucinlp/drop",
+            split="train",
+        ),
+    ),
+    "narrativeqa": dict(
+        train=dict(
+            path="nvidia/ChatQA-Training-Data",
+            name="narrativeqa",
+            split="train",
+        ),
+    ),
+    "quoref": dict(
+        train=dict(
+            path="nvidia/ChatQA-Training-Data",
+            name="quoref",
+            split="train",
+        ),
+    ),
+    "ropes": dict(
+        train=dict(
+            path="nvidia/ChatQA-Training-Data",
+            name="ropes",
+            split="train",
+        ),
+    ),
+    "synthetic_convqa": dict(
+        train=dict(
+            path="nvidia/ChatQA-Training-Data",
+            name="synthetic_convqa",
+            split="train",
+        ),
+    ),
+    "booksum": dict(
+        train=dict(
+            path="kmfoda/booksum",
+            split="train",
+        ),
+        validation=dict(
+            path="kmfoda/booksum",
+            split="validation",
+        ),
+        test=dict(
+            path="kmfoda/booksum",
+            split="test",
+        ),
+    ),
+    "gov_report": dict(
+        train=dict(
+            path="ccdv/govreport-summarization",
+            split="train",
+        ),
+        validation=dict(
+            path="ccdv/govreport-summarization",
+            split="validation",
+        ),
+        test=dict(
+            path="ccdv/govreport-summarization",
+            split="test",
+        ),
+    ),
+    # "wikitext-2": dict(
+    #     train=dict(
+    #         path="EleutherAI/wikitext_document_level",
+    #         name="wikitext-2-raw-v1",
+    #         split="train",
+    #     ),
+    # ),
+    # "wikitext-103": dict(
+    #     train=dict(
+    #         path="EleutherAI/wikitext_document_level",
+    #         name="wikitext-103-raw-v1",
+    #         split="train",
+    #     ),
+    # ),
+    "gsm8k": dict(
+        train=dict(
+            path="openai/gsm8k",
+            name="main",
+            split="train[100:]",
+        ),
+        validation=dict(
+            path="openai/gsm8k",
+            name="main",
+            split="train[:100]",
+        ),
+    ),
+    "openmathintx-2": dict(
+        train=dict(
+            path="nvidia/OpenMathInstruct-2",
+            split="train_1M[:100000]",
+        ),
+    ),
+    "opencoder-edu": dict(
+        train=dict(
+            path="OpenCoder-LLM/opc-sft-stage2",
+            name="educational_instruct",
+            split="train[900:]",
+        ),
+        validation=dict(
+            path="OpenCoder-LLM/opc-sft-stage2",
+            name="educational_instruct",
+            split="train[:900]",
+        ),
+    ),
+    "triviaqa_retrieved": dict(
+        test=dict(
+            path="json",
+            data_files="data/eval/triviaqa/triviaqa_retrieved.jsonl",
+            split="train",
+        ),
+    ),
+    # for negative training and evaluating
+    "negative_nq": dict(
+        test=dict(
+            path="json",
+            data_files="data/raw_datasets/negative_natural_questions/validation.jsonl",
+            split="train",
+        ),
+    ),
+}
+
+# LongBench kwargs
+for ds_name in LONGBENCH_TASKS + LONGBENCH_E_TASKS:
+    DS_KWARGS[ds_name] = dict(
+        test=dict(
+            path="THUDM/LongBench",
+            name=ds_name.split("/")[-1],
+            split="test",
+        )
+    )
+
+# for training closed qa datasets, e.g., hotpot_qa, squad, etc.
+CLOSED_QA_INTX_TEMPLATES = [
+    "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}",
+    "Answer without any explanation.\n\nQuestion: {input}",
+    "Based on the provided text, what is the answer to the following question? Provide only the answer.\n\nQuestion: {input}",
+    "Extract the answer to the question from the text. Be concise. Do not explain.\n\nQuestion: {input}",
+    "What is the answer to this question, based on the context? Respond with the answer only.\n\nQuestion: {input}",
+    "Provide a direct answer to the question using the given passages. Do not give any explanation.\n\nQuestion: {input}",
+    "Answer the question using only information from the provided text. No extra words.\n\nQuestion: {input}",
+    "From the passages, answer the question. Just the answer, please.\n\nQuestion: {input}",
+    "Give the answer to the question. Do not include any other text.\n\nQuestion: {input}",
+    "The answer to the question is in the text. Find it and state it clearly. No need for explanation.\n\nQuestion: {input}",
+    "Concisely answer the question based on the text provided. Don't include any other words. Just the answer.\n\nQuestion: {input}",
+    "Read the passages and answer the question with the minimal necessary words.\n\nQuestion: {input}",
+    "What is the direct response to the question, according to the text? Avoid explanation.\n\nQuestion: {input}",
+    "Please provide only the answer to the question, derived from the text.\n\nQuestion: {input}",
+    "Using the provided context, answer the question. Output the answer and nothing else.\n\nQuestion: {input}",
+    "Identify the answer in the text and present it without elaboration.\n\nQuestion: {input}",
+    "Answer the following question based on the text. Your answer should be brief and to the point. No explanation.\n\nQuestion: {input}",
+    "Based on the information given, what is the answer to the question? Only state the answer.\n\nQuestion: {input}",
+    "Find the answer to the question in the provided passages and write it down. No explanations.\n\nQuestion: {input}",
+    "The question is: {input}. Provide the answer based on the text, and nothing more.",
+    "Question: {input}\nAnswer directly based on the text provided. No extra words.",
+    "Question: {input}\nPlease provide the answer based on the text. No explanation is needed.",
+]
+
+
+EVAL_INTX_TEMPLATES = {
+    "negative_nq": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
+    "triviaqa_retrieved": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
+    "hotpot_qa": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
+    "squad": "Answer the following question. Output only the answer and do not output any other words.\n\nQuestion: {input}",
+    "longbench/qasper": 'Answer the question as concisely as you can, using a single phrase or sentence if possible.\nIf the question cannot be answered based on the information in the article, write "unanswerable".\nIf the question is a yes/no question, answer "yes", "no", or "unanswerable". Do not provide any explanation.\n\nQuestion: {input}',
+    "longbench/narrativeqa": "Answer the question as concisely as you can, using a single phrase if possible. Do not provide any explanation.\n\nQuestion: {input}",
+    "longbench/multifieldqa_en": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}",
+    # "longbench/multifieldqa_zh": "阅读以下文字并用中文简短回答：\n\n现在请基于上面的文章回答下面的问题，只告诉我答案，不要输出任何其他字词。\n\n问题：{input}",
+    "longbench/hotpotqa": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}",
+    "longbench/2wikimqa": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}",
+    "longbench/musique": "Answer the following question. Only output the answer and do not output any other words.\n\nQuestion: {input}",
+    #
+    # "longbench/dureader": "请基于给定的文章回答下述问题\n\n请基于上述文章回答下面的问题。\n\n问题：{input}",
+    "longbench/gov_report": "You are given a report by a government agency. Write a one-page summary of the report.",
+    "longbench/qmsum": "You are given a meeting transcript and a query containing a question or instruction. Answer the query based on the above meeting transcript.\n\nQuery: {input}",
+    "longbench/multi_news": "You are given several news passages. Write a one-page summary of all the news.",
+    # "longbench/vcsum": "下面有一段会议记录，请你阅读后，写一段总结，总结会议的内容。",
+}
+for ds_name in LONGBENCH_E_TASKS:
+    EVAL_INTX_TEMPLATES[ds_name] = EVAL_INTX_TEMPLATES[ds_name[:-2]]
+
+for ds_name in DS_KWARGS:
+    if ds_name not in EVAL_INTX_TEMPLATES:
+        EVAL_INTX_TEMPLATES[ds_name] = "{input}"

src/ctx_to_lora/data/pretrain.py

@@ -0,0 +1,5 @@
+# when generating dont forget to remove dup contexts
+# merge QAs from dup contexts to construct
+# a doc with QAs at the end
+# maybe mix with some that doesn;t have QAs
+# TODO: implement

src/ctx_to_lora/data/processing.py

@@ -0,0 +1,985 @@
+import hashlib
+import json
+import logging
+import random
+from collections.abc import Callable, Iterator
+from os import path
+from typing import Any
+
+import datasets
+import numpy as np
+from datasets import load_dataset
+from transformers import PreTrainedTokenizerBase
+
+from ctx_to_lora.data.definitions import (
+    CLOSED_QA_INTX_TEMPLATES,
+    DS_KWARGS,
+    EVAL_INTX_TEMPLATES,
+    IGNORE_INDEX,
+    TRANSFORMED_DATA_DIR,
+)
+from ctx_to_lora.utils import TRAINING_TASK
+
+logger = logging.getLogger()
+
+# TODO: add continued pre-training pipeline
+
+
+def closed_qa_prompting(prompt: str):
+    template = random.choice(CLOSED_QA_INTX_TEMPLATES)
+    return template.format(input=prompt)
+
+
+# TODO: dont use closed_qa_prompting when training on self-generated responses
+def get_preprocessing_fn(
+    ds_name: str, is_eval: bool
+) -> Callable[[dict[str, Any]], dict[str, Any]]:
+    """
+    Get preprocessing function for a specific dataset.
+
+    Args:
+        ds_name: Name of the dataset
+
+    Returns:
+        A preprocessing function that takes and returns a dictionary
+    """
+    f = lambda x: x
+
+    if ds_name.startswith("pwc"):
+
+        def f(sample):
+            return {
+                "context": sample["input"],
+                "prompt": sample["prompt"],
+                "response": sample["answer"],
+            }
+
+    elif ds_name.startswith("hotpot_qa"):
+
+        def f(sample):
+            txt = ""
+            for p in sample["context"]["sentences"]:
+                txt += " " + "".join(p)
+
+            q = sample["question"]
+            prompt = closed_qa_prompting(q) if not is_eval else q
+
+            return {
+                "context": txt.strip(),
+                "prompt": prompt,
+                "response": sample["answer"],
+            }
+
+    elif "squad" in ds_name:
+
+        def f(sample):
+            q = sample["question"]
+            prompt = closed_qa_prompting(q) if not is_eval else q
+            return {
+                "context": sample["context"],
+                "prompt": prompt,
+                "response": sample["answers"]["text"][0],
+            }
+
+    elif ds_name == "drop":
+
+        def f(sample):
+            q = sample["question"]
+            prompt = closed_qa_prompting(q) if not is_eval else q
+            return {
+                "context": sample["passage"],
+                "prompt": prompt,
+                "response": ", ".join(sample["answers_spans"]["spans"]),
+            }
+
+    elif ds_name in ["narrativeqa", "quoref", "ropes"]:
+
+        def f(sample):
+            response = sample["answers"][0]
+            if isinstance(response, list):
+                response = response[0]
+            q = sample["messages"][-1]["content"]
+            prompt = closed_qa_prompting(q) if not is_eval else q
+            return {
+                "context": sample["document"],
+                "prompt": prompt,
+                "response": response,
+            }
+
+    elif ds_name == "synthetic_convqa":
+
+        def f(sample):
+            response = sample["answers"][0]
+            if isinstance(response, list):
+                response = response[0]
+            return {
+                "context": sample["document"],
+                "prompt": sample["messages"][-1]["content"],
+                "response": response,
+            }
+
+    elif ds_name == "booksum":
+        prompt_templates = [
+            "Summarization the provided text.",
+            "# Summary",
+            "### Summary",
+            "Summary of the text",
+        ]
+        # TODO: use these templates for training
+        # analysis_templates = []
+        # summary_templates = []
+
+        def f(sample):
+            return {
+                "context": sample["chapter"],
+                "prompt": "Summarize the provided text.",
+                "response": sample["summary_text"],
+            }
+
+    elif ds_name == "gov_report":
+
+        def f(sample):
+            return {
+                "context": sample["report"],
+                "prompt": "Summarize the provided text.",
+                "response": sample["summary"],
+            }
+
+    elif "wikitext" in ds_name:
+
+        def f(sample):
+            return {
+                "context": sample["page"],
+                "prompt": "PLAECHOLDER",
+                "response": "PLAECHOLDER",
+            }
+
+    elif ds_name == "openmathintx-2":
+
+        def f(sample):
+            return {
+                "context": sample["problem"],
+                "prompt": sample["problem"],
+                "response": sample["generated_solution"],
+            }
+
+    elif "gsm8k" in ds_name:
+
+        def f(sample):
+            return {
+                "context": sample["question"],
+                "prompt": sample["question"],
+                "response": sample["answer"],
+            }
+
+    elif "opencoder-edu" in ds_name:
+
+        def f(sample):
+            return {
+                "context": sample["instruction"],
+                "prompt": sample["instruction"],
+                "response": "```python\n" + sample["code"].strip() + "\n```",
+            }
+
+    elif ds_name.startswith("longbench"):
+
+        def f(sample):
+            return {
+                "context": sample["context"],
+                "prompt": sample["input"],
+                "response": sample["answers"][0],
+            }
+
+    elif ds_name == "negative_nq":
+
+        def f(sample):
+            q = sample["prompt"]
+            prompt = closed_qa_prompting(q) if not is_eval else q
+            return {
+                "context": sample["context"],
+                "prompt": prompt,
+                "response": sample["answer"],
+            }
+
+    elif ds_name == "triviaqa_retrieved":
+        # only used for eval
+        def f(sample):
+            return {
+                "context": sample["context"],
+                "prompt": sample["prompt"],
+                "response": sample["answer"],
+            }
+
+    if is_eval and ds_name in EVAL_INTX_TEMPLATES:
+        prompt_template = EVAL_INTX_TEMPLATES[ds_name]
+
+        def eval_intx_decorator(f):
+            def g(sample):
+                sample = f(sample)
+                sample["prompt"] = prompt_template.format(input=sample["prompt"])
+                return sample
+
+            return g
+
+        f = eval_intx_decorator(f)
+
+    return f
+
+
+def get_ds_kwargs(ds_name: str, split: str) -> dict[str, Any]:
+    if (ds_name not in DS_KWARGS) or (split not in DS_KWARGS[ds_name]):
+        kwargs = dict(path=ds_name, split=split)
+        logger.warning(
+            f"No dataset kwargs found for '{ds_name}' with split '{split}'.\n"
+            f"Using default kwargs: {kwargs}"
+        )
+        # return kwargs
+    else:
+        # return DS_KWARGS[ds_name][split]
+        kwargs = DS_KWARGS[ds_name][split]
+
+    # custom logic for slicing iterable datasets
+    take, skip = None, None
+    kw_split = kwargs["split"]
+    if ("[" in kw_split) and kw_split.endswith("]"):
+        kwargs["split"], slice = kw_split.split("[")
+        slice = slice.strip("]")
+        skip = slice.split(":")[0]
+        if skip:
+            kwargs["skip"] = int(skip)
+        take = slice.split(":")[1]
+        if take:
+            kwargs["take"] = int(take)
+    return kwargs
+
+
+def validate_columns(tokenized_ds):
+    cols = ["input_ids", "attention_mask", "labels"]
+    if "ctx_ids" in tokenized_ds.column_names:
+        cols += ["ctx_ids", "ctx_attn_mask"]
+    if "chat_ids" in tokenized_ds.column_names:
+        cols += ["chat_ids", "chat_attn_mask", "chat_labels"]
+    ref_cols = set(cols)
+    assert set(tokenized_ds.column_names) == ref_cols, (
+        f"Columns mismatch: {set(tokenized_ds.column_names)} != {ref_cols}"
+    )
+
+
+def filter_long_samples(samples):
+    return [len(ctx) < 10000 for ctx in samples["context"]]
+
+
+def filter_long_chat(samples):
+    return [len(chat) < 2**13 for chat in samples["chat"]]
+
+
+def add_repeat_prompt_fn(samples):
+    unique_contexts = set()
+    ctxs, prompts, responses = [], [], []
+    for ctx, prompt, response in zip(
+        samples["context"], samples["prompt"], samples["response"]
+    ):
+        # Only process if the context is not already in the set
+        if ctx in unique_contexts:
+            continue
+        # remove too long contexts
+        if len(ctx) > 1000:
+            continue
+        unique_contexts.add(ctx)
+        ctxs.append(ctx)
+        responses.append(ctx)
+        prompts.append("Repeat the text above.")
+
+    logger.debug("Adding repeat prompt...")
+    logger.debug(f"# unique contexts: {len(unique_contexts)}")
+
+    return dict(
+        context=ctxs + samples["context"],
+        response=responses + samples["response"],
+        prompt=prompts + samples["prompt"],
+    )
+
+
+def add_negative_prompt_fn(samples):
+    unique_contexts = set()
+    ctxs, prompts, responses = [], [], []
+    keywords = [
+        "repeat",
+        "rephrase",
+        "summarize",
+        "rewrite",
+        "title",
+        "keyword",
+        "continuation",
+    ]
+
+    for ctx, prompt, response in zip(
+        samples["context"], samples["prompt"], samples["response"]
+    ):
+        if ctx in unique_contexts:
+            continue
+        unique_contexts.add(ctx)
+        if len(ctx) > 3000:
+            continue
+        if any(keyword in prompt for keyword in keywords):
+            # Skip samples where the prompt contains any of the specified keywords
+            continue
+
+        ctxs.append(ctx)
+        prompts.append(prompt)
+        responses.append(response)
+
+    logger.debug("Adding negative prompt...")
+    logger.debug(f"# unique contexts: {len(unique_contexts)}")
+
+    # remove one last sample if the number of samples is odd
+    if len(ctxs) % 2 != 0:
+        ctxs.pop()
+        prompts.pop()
+        responses.pop()
+
+    # to make sure that the negative prompt/response is not the same as the original
+    indices = list(np.random.permutation(len(ctxs))) + list(
+        np.random.permutation(len(ctxs))
+    )
+    neg_ctxs, neg_prompts, neg_responses = [], [], []
+    for idx in range(0, len(indices), 2):
+        i = indices[idx]
+        j = indices[idx + 1]
+        neg_ctxs.append(ctxs[i])
+        neg_prompts.append(ctxs[j] + "\n\n" + prompts[j])
+        neg_responses.append(responses[j])
+
+    return dict(
+        context=neg_ctxs + samples["context"],
+        prompt=neg_prompts + samples["prompt"],
+        response=neg_responses + samples["response"],
+    )
+
+
+def filter_none(samples):
+    out = [True] * len(samples["context"])
+    for k in samples:
+        for i, sample in enumerate(samples[k]):
+            if not bool(sample):
+                out[i] = False
+    return out
+
+
+def _load_and_process_dataset(
+    ds_name: str,
+    split: str,
+    add_negative_prompt: bool,
+    add_repeat_prompt: bool,
+    streaming: bool,
+    ds_path: str,
+    num_proc: int,
+):
+    try:
+        ds_kwargs = get_ds_kwargs(ds_name, split)
+        skip = ds_kwargs.pop("skip", None)
+        take = ds_kwargs.pop("take", None)
+        # only load as stream if the dataset is local (parquet)
+        load_as_stream = (
+            ("fw_qa" in ds_name)
+            and ("tiny" not in ds_name)
+            and ("mini" not in ds_name)
+            and streaming
+            and ds_kwargs["path"] == "parquet"
+        )
+        ds = load_dataset(
+            **ds_kwargs,
+            trust_remote_code=True,
+            streaming=load_as_stream,
+        )
+        if split == "train" and streaming and (not load_as_stream):
+            # if the dataset is hosted on HF, we load it first then convert to iterable
+            ds = ds.to_iterable_dataset(num_shards=128)
+        if skip is not None:
+            ds = ds.skip(skip)
+        if take is not None:
+            ds = ds.take(take)
+    except ValueError as e:
+        logger.warning(
+            f"Failed to load dataset {ds_name} with split {split}. Error: {e}\nSkipping..."
+        )
+        return None
+    cols_to_remove = [
+        col for col in ds.column_names if col not in ["context", "prompt", "response"]
+    ]
+    is_eval = split != "train"
+    ds = ds.map(
+        get_preprocessing_fn(ds_name, is_eval),
+        remove_columns=cols_to_remove,
+        num_proc=num_proc,
+    )
+    # ds = ds.remove_columns(cols_to_remove)
+    ds = ds.filter(filter_none, batched=True, num_proc=num_proc)
+
+    if split == "train":
+        # TODO: drop samples or split into multiple contexts
+        ds = ds.filter(filter_long_samples, batched=True, num_proc=num_proc)
+        if add_negative_prompt:
+            # TODO: make explicit dataset for this
+            ds = ds.map(
+                add_negative_prompt_fn,
+                batched=True,
+                batch_size=100_000,
+                num_proc=num_proc,
+            )
+        if add_repeat_prompt and "context_numbers" not in ds_name:
+            # TODO: convert to AE/LM
+            ds = ds.map(
+                add_repeat_prompt_fn,
+                batched=True,
+                batch_size=100_000,
+                num_proc=num_proc,
+            )
+        ds.save_to_disk(ds_path, num_proc=num_proc)
+    return ds
+
+
+def get_tokenized_dataset(
+    ds_name: str,
+    split: str,
+    base_model_max_len: int,
+    tokenizer: PreTrainedTokenizerBase,
+    tokenizer_kwargs: dict[str, Any],
+    ctx_model_max_len: int,
+    ctx_tokenizer: PreTrainedTokenizerBase,
+    ctx_tokenizer_kwargs: dict[str, Any],
+    add_ctx_to_chat: bool,
+    add_repeat_prompt: bool,
+    add_negative_prompt: bool,
+    use_kl_loss: bool,
+    set_format: str | None = None,
+    streaming: bool = False,
+) -> dict[str, Any]:
+    # TODO: update hash function to include max_len
+    # TODO: max_len could be included in the kwargs
+    assert not use_kl_loss, "KL loss is deprecated"
+    logger.debug(f"Loading dataset {ds_name} with split {split}...")
+    need_ctx_ids = not add_ctx_to_chat and bool(ctx_model_max_len)
+
+    load_and_process_kwargs = dict(
+        ds_name=ds_name,
+        split=split,
+        add_negative_prompt=add_negative_prompt,
+        add_repeat_prompt=add_repeat_prompt,
+        streaming=streaming,
+    )
+    num_proc = None if streaming and split == "train" else 8
+
+    ds_hash = hashlib.sha256(json.dumps(load_and_process_kwargs).encode()).hexdigest()
+    ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}"
+
+    if path.exists(ds_path) and split == "train":
+        # load the cached ds
+        logger.info(f"Loaded processed dataset from {ds_path}")
+    else:
+        logger.info(f"Loading dataset {ds_name} with split {split}...")
+        ds = _load_and_process_dataset(
+            **load_and_process_kwargs,
+            ds_path=ds_path,
+            num_proc=num_proc,
+        )
+    if split == "train":
+        # force loading the cached version for newly created ds
+        ds = datasets.load_from_disk(ds_path)
+
+    tokenized_ds = construct_and_tokenize_ctx_qa(
+        base_model_max_len,
+        tokenizer,
+        tokenizer_kwargs,
+        ctx_model_max_len,
+        ctx_tokenizer,
+        ctx_tokenizer_kwargs,
+        add_ctx_to_chat,
+        use_kl_loss,
+        need_ctx_ids,
+        ds,
+        split,
+        set_format,
+        num_proc,
+    )
+    return tokenized_ds
+
+
+def construct_and_tokenize_ctx_qa(
+    base_model_max_len,
+    tokenizer,
+    tokenizer_kwargs,
+    ctx_model_max_len,
+    ctx_tokenizer,
+    ctx_tokenizer_kwargs,
+    add_ctx_to_chat,
+    use_kl_loss,
+    need_ctx_ids,
+    ds,
+    split,
+    set_format=None,
+    num_proc=None,
+):
+    is_eval = split != "train"
+    # TODO: update hash function to include max_len
+    # TODO: max_len could be included in the kwargs
+    kwargs = dict(
+        tokenizer=repr(tokenizer),
+        tokenizer_kwargs=json.dumps(tokenizer_kwargs),
+        ctx_tokenizer=repr(ctx_tokenizer),
+        ctx_tokenizer_kwargs=json.dumps(ctx_tokenizer_kwargs),
+        add_ctx_to_chat=add_ctx_to_chat,
+        use_kl_loss=use_kl_loss,
+        need_ctx_ids=need_ctx_ids,
+        ds=ds._fingerprint,
+        # TODO: add split (after moving to continued pre-training pipeline)
+        set_format=set_format,
+    )
+    kwargs_str = json.dumps(kwargs)
+    logger.debug(f"Tokenizing dataset with kwargs: {kwargs_str}")
+    ds_hash = hashlib.sha256(kwargs_str.encode()).hexdigest()
+    ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}"
+    if path.exists(ds_path) and split == "train":
+        # load the cached ds
+        logger.info(f"Loaded tokenized dataset from {ds_path}")
+        ds = datasets.load_from_disk(ds_path)
+        return ds
+    # for sft + chat_model, we need to convert the dataset to chat format
+    # add "messages" field
+    ds = ds.map(
+        convert_ctx_prompt_response_to_messages,
+        fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat},
+        num_proc=num_proc,
+    )
+    # add "chat" field
+    ds = ds.map(
+        get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer),
+        batched=True,
+        batch_size=100_000,
+        num_proc=num_proc,
+    )
+    # tokenize the chat + mask the assistant inputs
+    # TODO: drop samples or split into multiple contexts
+    if split == "train":
+        ds = ds.filter(
+            filter_long_chat,
+            batched=True,
+            batch_size=100_000,
+            num_proc=num_proc,
+        )
+
+    # add "input_ids", "attention_mask", "labels"
+    tokenized_ds = ds.map(
+        tokenize_chat_messages,
+        fn_kwargs={
+            "tokenizer": tokenizer,
+            "mask_assistant_inputs": True,
+            "tokenizer_kwargs": tokenizer_kwargs,
+        },
+        num_proc=num_proc,
+    )
+    if split == "train":
+        # drop?
+        # base model cant handle these samples naturally
+        # should skip
+        ...
+    else:
+        tokenized_ds = tokenized_ds.map(
+            truncate_middle_if_too_long,
+            fn_kwargs={
+                "max_length": base_model_max_len,
+                "columns": ["input_ids", "attention_mask"],
+                "max_new_tokens": 256,
+            },
+            batched=True,
+            batch_size=100_000,
+            num_proc=num_proc,
+        )
+        tokenized_ds = tokenized_ds.map(
+            add_length_info,
+            fn_kwargs={"columns": ["input_ids"]},
+            batched=True,
+            batch_size=100_000,
+            num_proc=num_proc,
+        )
+
+    # for use_kl_loss, we need "chat_ids" and "chat_attn_mask"
+    if use_kl_loss:
+        raise NotImplementedError("KL loss deprecated")
+        tokenized_ds = tokenized_ds.map(
+            convert_ctx_prompt_response_to_messages,
+            fn_kwargs={"add_ctx_to_chat": True},
+            remove_columns=["messages"],
+        )
+        tokenized_ds = tokenized_ds.map(
+            get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer),
+            remove_columns=["chat"],
+        )
+        tokenized_ds = tokenized_ds.map(
+            tokenize_chat_messages,
+            fn_kwargs={
+                "tokenizer": tokenizer,
+                "mask_assistant_inputs": True,
+                "tokenizer_kwargs": tokenizer_kwargs,
+                "for_kl_loss": True,
+            },
+        )
+
+    if need_ctx_ids:
+        # tokenize the ctx_text to get ctx_ids and ctx_attn_mask
+        tokenized_ds = tokenized_ds.map(
+            tokenize_ctx_text,
+            fn_kwargs={"tokenizer": ctx_tokenizer},
+            batched=True,
+            batch_size=100_000,
+            num_proc=num_proc,
+        )
+        if split == "train":
+            # TODO: do something if ctx length is longer than the ctx model length
+            # e.g., drop or split for multi-lora training
+            ...
+
+        else:
+            # truncate in the middle
+            tokenized_ds = tokenized_ds.map(
+                truncate_middle_if_too_long,
+                fn_kwargs={
+                    "max_length": ctx_model_max_len,
+                    "columns": ["ctx_ids", "ctx_attn_mask"],
+                    # no need to add new_tokens here (used only for the ctx encoder)
+                    "max_new_tokens": 0,
+                },
+                batched=True,
+                batch_size=100_000,
+                num_proc=num_proc,
+            )
+            tokenized_ds = tokenized_ds.map(
+                add_length_info,
+                fn_kwargs={"columns": ["ctx_ids"]},
+                batched=True,
+                batch_size=100_000,
+                num_proc=num_proc,
+            )
+    tokenized_ds = tokenized_ds.remove_columns(
+        ["messages", "chat", "context", "prompt", "response"]
+    )
+    if set_format:
+        tokenized_ds.set_format(type=set_format)
+    # if isinstance(tokenized_ds, IterableDataset):
+    #     # the columns are unknown when using streaming dataset
+    #     tokenized_ds = tokenized_ds._resolve_features()
+    # validate_columns(tokenized_ds)
+    if split == "train":
+        tokenized_ds.save_to_disk(ds_path, num_proc=num_proc)
+        del tokenized_ds
+        return datasets.load_from_disk(ds_path)
+    return tokenized_ds
+
+
+def get_sft_prompt_formatting_fn(
+    sft_mode: TRAINING_TASK,
+    tokenizer: PreTrainedTokenizerBase,
+) -> Callable[[dict[str, Any]], dict[str, Any]]:
+    """
+    Get a function that formats examples for supervised fine-tuning.
+
+    Args:
+        sft_mode: The training task type
+        tokenizer: The tokenizer to use for chat template application
+
+    Returns:
+        A function that takes a training example and returns formatted data
+
+    Raises:
+        NotImplementedError: If sft_mode is not COMPLETION or tokenizer has no chat template
+    """
+    if sft_mode != TRAINING_TASK.COMPLETION:
+        raise NotImplementedError(
+            f"Training task {sft_mode} not supported. "
+            "Only completion is supported for now."
+        )
+
+    if tokenizer.chat_template is None:
+        raise NotImplementedError(
+            "Only chat models + SFT are supported. "
+            "Training with pre-training data is not supported yet."
+        )
+    # TODO: add support for causal_lm training (for pre-training data)
+    # TODO: add support for recon training (for pre-training data)
+    # TODO: add support for non-chat models
+
+    # def f(example):
+    #     output_texts = (
+    #         dict(text=[]) if sft_mode == "causal_lm" else dict(prompt=[], response=[])
+    #     )
+    #     df = pd.DataFrame(dict(example))
+    #     for i, inp_txt in df.iterrows():
+    #         if sft_mode == "causal_lm":
+    #             text = metadata["text_template"].format(**inp_txt)
+    #             output_texts["text"].append(text)
+    #         elif sft_mode == "completion":
+    #             prompt = metadata["user_prompt_template"].format(**inp_txt)
+    #             output_texts["prompt"].append(prompt)
+    #             output_texts["response"].append(str(inp_txt[metadata["response_field"]]))
+    #     return output_texts
+
+    def f_intx(example):
+        chat_text = tokenizer.apply_chat_template(
+            example["messages"], tokenize=False, add_generation_prompt=False
+        )
+        return dict(chat=chat_text)
+
+    # return f if not apply_chat_template_fn is not None else f_intx
+    return f_intx
+
+
+def convert_ctx_prompt_response_to_messages(
+    example: dict[str, Any],
+    add_ctx_to_chat: bool = True,
+) -> dict[str, Any]:
+    """
+    Convert context/prompt/response format to chat messages format.
+
+    Args:
+        example: Dictionary containing 'prompt' and 'response' keys
+        add_ctx_to_chat: Whether to prepend context to the user message
+
+    Returns:
+        Dictionary with added 'messages' key containing chat format
+
+    Raises:
+        ValueError: If 'prompt' or 'response' keys are missing
+    """
+    if "prompt" not in example or "response" not in example:
+        raise ValueError(f"'prompt' and 'response' are required. Got: {example}")
+
+    system_msg = ""
+    if "system_message" in example:
+        system_msg = example["system_message"].strip()
+
+    user_msg = example["prompt"].strip()
+    if add_ctx_to_chat:
+        user_msg = example["context"].strip() + "\n\n" + user_msg.strip()
+
+    messages = [
+        {"role": "system", "content": system_msg.strip()},
+        {"role": "user", "content": user_msg.strip()},
+        {"role": "assistant", "content": example["response"].strip()},
+    ]
+
+    return dict(messages=messages)
+
+
+# adapted from https://github.com/huggingface/trl/issues/632#issuecomment-1972630547
+def get_assistant_start_end_indices(
+    messages: list[dict[str, str]],
+    tokenizer: PreTrainedTokenizerBase,
+    conversation_ids: list[int],
+) -> list[tuple[int, int]]:
+    """
+    Get the start and end indices of assistant messages in conversation text using tokenized messages.
+
+    Args:
+        messages: List of message dictionaries with 'role' and 'content' keys
+        tokenizer: Tokenizer used for tokenization
+        conversation_ids: Tokenized conversation
+
+    Returns:
+        List of (start, end) index tuples for assistant messages in the tokenized conversation
+    """
+    assistant_ranges = []
+    current_idx = 0
+
+    for idx, message in enumerate(messages):
+        # if message["role"] == "assistant":
+        # Tokenize the assistant message content
+        tokenized_message = tokenizer.encode(
+            message["content"], add_special_tokens=False
+        )
+        # Find the start index of the assistant message in the tokenized conversation
+        for i in range(current_idx, len(conversation_ids)):
+            if conversation_ids[i : i + len(tokenized_message)] == tokenized_message:
+                start_idx = i
+                end_idx = i + len(tokenized_message)
+                current_idx = end_idx
+                if message["role"] == "assistant":
+                    if idx == (len(messages) - 1):
+                        # include everything after the last assistant message
+                        assistant_ranges.append((start_idx, len(conversation_ids)))
+                    else:
+                        assistant_ranges.append((start_idx, end_idx))
+                break
+
+    return assistant_ranges
+
+
+def get_masked_labels(
+    conversation_ids: dict[str, list[Any]],
+    assistant_ranges: list[tuple[int, int]],
+) -> Iterator[int]:
+    """
+    Generate masked labels for conversation, masking non-assistant tokens.
+    NOTE: This will also includes extra tokens between assistant and user messages
+    in multi-turn conversations.
+    E.g., {assistant_msg} <|eot_id|><|start_header_id|>user<|end_header_id|> {user_msg}
+    for Llama 3 models.
+
+    Args:
+        conversation_ids: Dictionary with tokenized conversation info
+        assistant_ranges: List of (start, end) indices for assistant messages
+
+    Yields:
+        Token ID or IGNORE_INDEX for each position
+    """
+    for id_, (id_s, id_e) in list(
+        zip(conversation_ids["input_ids"], conversation_ids["offset_mapping"])
+    ):
+        if any(id_s >= s and id_e <= e for s, e in assistant_ranges):
+            yield id_
+        else:
+            yield IGNORE_INDEX
+
+
+def tokenize_chat_messages(
+    example: dict[str, Any],
+    tokenizer: PreTrainedTokenizerBase,
+    mask_assistant_inputs: bool = True,
+    for_kl_loss: bool = False,
+    tokenizer_kwargs: dict[str, Any] | None = None,
+) -> dict[str, list[int]]:
+    """
+    Tokenize chat messages and optionally mask non-assistant tokens.
+
+    Args:
+        example: Dictionary containing 'chat' and 'messages' keys
+        tokenizer: Tokenizer to use
+        mask_assistant_inputs: Whether to mask non-assistant tokens
+        for_kl_loss: change the column names to "chat_ids" and "chat_attn_mask"
+        tokenizer_kwargs: Additional arguments to pass to tokenizer
+
+    Returns:
+        Dictionary containing tokenized inputs and labels
+    """
+    # should be used only with chat models
+    text = example["chat"]
+    messages = example["messages"]
+    n_response = len([m for m in messages if m["role"] == "assistant"])
+    if n_response != 1:
+        raise ValueError(f"Expected 1 assistant response. Got {n_response}.")
+    conversation_ids = tokenizer(
+        text,
+        # return_offsets_mapping=mask_assistant_inputs,
+        add_special_tokens=False,
+        truncation=False,
+        **(tokenizer_kwargs or {}),
+    )
+    if tokenizer_kwargs and (
+        len(conversation_ids["input_ids"]) >= tokenizer_kwargs["max_length"]
+    ):
+        raise ValueError(
+            f"Conversation length {len(conversation_ids['input_ids'])} exceeds max length {tokenizer_kwargs['max_length']}"
+        )
+
+    if mask_assistant_inputs:
+        assistant_ranges = get_assistant_start_end_indices(
+            messages, tokenizer, conversation_ids["input_ids"]
+        )
+        # labels = get_masked_labels(conversation_ids, assistant_ranges)
+        labels = [
+            id_ if any(s <= i < e for s, e in assistant_ranges) else IGNORE_INDEX
+            for i, id_ in enumerate(conversation_ids["input_ids"])
+        ]
+        conversation_ids["labels"] = labels
+        # del conversation_ids["offset_mapping"]
+    else:
+        conversation_ids["labels"] = conversation_ids["input_ids"]
+    if for_kl_loss:
+        conversation_ids["chat_ids"] = conversation_ids.pop("input_ids")
+        conversation_ids["chat_attn_mask"] = conversation_ids.pop("attention_mask")
+        conversation_ids["chat_labels"] = conversation_ids.pop("labels")
+    return conversation_ids
+
+
+def add_length_info(samples: dict[str, any], columns: list[str]) -> dict[str, any]:
+    for col in columns:
+        samples[f"{col}_len"] = [len(t) for t in samples[col]]
+    return samples
+
+
+def truncate_middle_if_too_long(
+    samples: dict[str, any],
+    max_length: int,
+    columns: list[str],
+    max_new_tokens: int = 256,
+) -> dict[str, any]:
+    """
+    Truncate the middle of a list of tokens to fit within a maximum length.
+
+    Args:
+        tokens: List of token IDs
+        max_length: Maximum length for the truncated tokens
+
+    Returns:
+        List of truncated token IDs
+    """
+    max_new_tokens_half = max_new_tokens // 2
+    # leave max_new_tokens for generation
+    half = max_length // 2 - max_new_tokens_half
+    for col in columns:
+        # if len(samples[col]) <= max_length:
+        #     return samples[col]
+        samples[col] = [
+            t[:half] + t[-half:] if len(t) > max_length else t for t in samples[col]
+        ]
+    return samples
+
+
+def tokenize_ctx_text(
+    samples: dict[str, Any],
+    tokenizer: PreTrainedTokenizerBase,
+) -> dict[str, Any]:
+    if tokenizer.chat_template:
+        txt = tokenizer.apply_chat_template(
+            [
+                [{"role": "user", "content": text.strip()}]
+                for text in samples["context"]
+            ],
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+        # special tokens are already added by the chat template
+        tokenized_text = tokenizer(txt, truncation=False, add_special_tokens=False)
+    else:
+        txt = samples["context"]
+        tokenized_text = tokenizer(txt, truncation=False)
+
+    ctx_ids = tokenized_text["input_ids"]
+    ctx_attn_mask = tokenized_text["attention_mask"]
+    return dict(ctx_ids=ctx_ids, ctx_attn_mask=ctx_attn_mask)
+
+
+if __name__ == "__main__":
+    from transformers import AutoTokenizer
+
+    # model_name = "meta-llama/Llama-3.2-1B-Instruct"
+    # model_name = "meta-llama/Llama-2-7b-chat-hf"
+    model_name = "google/gemma-2-2b-it"
+    messages = [
+        {"role": "user", "content": "Hello!"},
+        {"role": "assistant", "content": "Hello!"},
+        # {"role": "user", "content": "Not too bad"},
+        # {"role": "assistant", "content": "Cooooooool"},
+    ]
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    chat = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=False, add_special_tokens=False
+    )
+    print(chat)
+    model_inputs = tokenize_chat_messages(
+        {"chat": chat, "messages": messages},
+        tokenizer,
+        # for_kl_loss=True,
+    )
+
+    print(tokenizer(chat, add_special_tokens=False))
+    print(model_inputs)