如何将Dataset的label列转为ClassLabel以支持train_test_split分层？

步骤1：将Sequence类型的label转为单值 #

如果你的label列是单元素序列（比如每个样本的label是[0]、[1]这类单元素列表），先提取出列表中的唯一值：

ds = ds.map(lambda x: {"label": x["label"][0]})

如果是多标签序列，需根据业务逻辑将其转换为单个类别（比如取第一个标签、取出现频次最高的标签等）。

步骤2：将label列转换为ClassLabel类型 #

使用datasets.ClassLabel来转换列类型：

from datasets import ClassLabel

# 获取所有唯一标签值，生成ClassLabel对象
unique_labels = ds.unique("label")
class_label_type = ClassLabel(num_classes=len(unique_labels), names=unique_labels)

# 转换label列的类型
ds = ds.cast_column("label", class_label_type)

步骤3：执行分层拆分 #

现在可以正常使用train_test_split的分层参数：

train_testvalid = ds.train_test_split(test_size=0.5, shuffle=True, stratify_by_column="label")

补充说明 #

如果label是真正的多标签序列（一个样本对应多个类别），直接转ClassLabel不适用，因为分层拆分仅支持单标签任务。这种情况需重新确定分层依据：比如选一个核心标签作为分层列，或者将多标签组合为唯一字符串作为新的单标签列。

内容的提问来源于stack exchange，提问作者Yana