import pandas as pd
from datasets import load_dataset


def prepare_data():
    dataset = load_dataset("missvector/linux-commands", split="train")

    data_list = []
    for item in dataset:
        # 构建符合 ChatML 的对话格式
        sample = {
            "text": f"<|im_start|>user\nConvert to Linux command: {item['ch']}<|im_end|>\n<|im_start|>assistant\n{item['completion']}<|im_end|>"
        }
        data_list.append(sample)

    # 划分训练集和验证集
    df = pd.DataFrame(data_list)
    train_df = df.sample(frac=0.95, random_state=42)
    valid_df = df.drop(train_df.index)

    # 保存为 jsonl
    train_df.to_json("train.jsonl", orient="records", lines=True, force_ascii=False)
    valid_df.to_json("valid.jsonl", orient="records", lines=True, force_ascii=False)
    print("Data prepared: train.jsonl and valid.jsonl")


if __name__ == "__main__":
    prepare_data()