import json

from datasets import Sample, load_dataset  # type: ignore

dataset = load_dataset("json", data_files="dataset_2.jsonl", split="train")


def validate_roles(sample: Sample) -> bool:
    return all(msg["role"] in ("user", "assistant") for msg in sample["messages"])


cleaned_dataset = dataset.filter(validate_roles)
cleaned_dataset.save_to_disk("cleaned_dataset")
with open("cleaned_dataset.jsonl", "w") as f:
    for example in cleaned_dataset:
        f.write(json.dumps(example) + "\n")
