"""
该脚本用于处理和转换数学问题数据集，包括数据加载、分割、模板应用和保存功能
"""
from datasets import load_dataset, Dataset, DatasetDict
import pandas as pd
from sklearn.model_selection import train_test_split
import os
from random import randint

# 数据集名称
DATASET_NAME = "ServiceNow-AI/R1-Distill-SFT"
# 数据集版本
DATASET_VERSION = "v0"

# 加载原始数据集
dataset = load_dataset(DATASET_NAME, DATASET_VERSION)
# 转换为DataFrame并取样
df = pd.DataFrame(dataset['train'])
# 取前500条据用于示例
df = df.iloc[0:500]
# 将数据集分割为训练集和测试集
train, test = train_test_split(df, test_size=0.1, random_state=666)

print("训练集样本数量: ", len(train))
print("测试集样本数量: ", len(test))

# 数学问题的提示模板
prompt_template = """
You are a mathematics expert, skilled in solving mathematical problems and logical reasoning.
Please read the question carefully and provide detailed solution steps.

<|Question|>
{problem}

<|Solution Process|>
{reannotated_assistant_content}

<|Final Answer|>
{solution}
"""

def template_dataset(sample):
    """
    将原始样本应用模板格式化
    """
    sample["text"] = prompt_template.format(
        problem=sample["problem"],
        reannotated_assistant_content=sample["reannotated_assistant_content"],
        solution=sample["solution"]
    )
    return sample

# 分别为训练集和测试集创建Dataset对象，并组合为DatasetDict
train_dataset = Dataset.from_pandas(train)
test_dataset = Dataset.from_pandas(test)

dataset = DatasetDict({
    "train": train_dataset, 
    "test": test_dataset
})

# 对训练集和测试集应用模板转换，使用map函数处理每个样本，并移除原始特征列
train_dataset = dataset["train"].map(
    template_dataset, 
    remove_columns=list(dataset["train"].features)
)

test_dataset = dataset["test"].map(
    template_dataset, 
    remove_columns=list(dataset["test"].features)
)

# 随机打印一个训练样本
print("\n示例数据:")
print(train_dataset[randint(0, len(train_dataset)-1)]["text"])

# 将处理后的数据集保存为JSONL格式
os.makedirs("data", exist_ok=True)
train_dataset.to_json("data/train.jsonl")
test_dataset.to_json("data/eval.jsonl")
