Felladrin/ChatML-databricks-dolly-15k

from datasets import load_dataset
import pandas
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(
    pretrained_model_name_or_path="Felladrin/Llama-160M-Chat-v1"
)

dataset = load_dataset("databricks/databricks-dolly-15k", split="train")


def format(columns):
    instruction = columns["instruction"].strip()
    context = columns["context"].strip()
    response = columns["response"].strip()

    if context:
        user_message = f"{instruction}\n\nContext:\n{context}"
    else:
        user_message = instruction

    messages = [
        {
            "role": "user",
            "content": user_message,
        },
        {
            "role": "assistant",
            "content": response,
        },
    ]

    return tokenizer.apply_chat_template(messages, tokenize=False)


pandas.DataFrame({"text": [format(columns) for columns in dataset]}).to_parquet("train.parquet", index=False)