如何在Hugging Face中从CSV文件加载自定义数据集

加载CSV格式的自定义数据集通常配合Hugging Face的datasets库实现，再和transformers库搭配完成后续的模型训练、推理流程，具体操作如下：

• 安装需要的第三方库，命令如下：
  pip install datasets transformers pandas

2.1 单CSV文件加载 #

如果你的所有数据都存在一个CSV文件里，直接调用load_dataset指定格式为csv即可：

from datasets import load_dataset

# 替换为你自己的CSV文件路径
dataset = load_dataset("csv", data_files="./custom_data.csv")

加载完成后得到的是DatasetDict结构，默认所有数据会被分到train分组下，你可以通过以下方式读取数据：

# 查看第一条数据
print(dataset["train"][0])
# 查看数据集总条数
print(len(dataset["train"]))

2.2 拆分数据集 #

已提前拆分训练/验证/测试集 #

如果你的数据已经拆成了多个CSV文件，给data_files参数传入字典指定不同分组的路径即可：

dataset = load_dataset(
    "csv",
    data_files={
        "train": "./train.csv",
        "validation": "./val.csv",
        "test": "./test.csv"
    }
)

单文件自动拆分 #

如果只有一个CSV文件需要拆分训练集和测试集，调用内置的train_test_split方法即可：

# 先加载全量数据
dataset = load_dataset("csv", data_files="./custom_data.csv")
# 按8:2拆分训练集和测试集，test_size可自行调整
dataset = dataset["train"].train_test_split(test_size=0.2)

如果需要做分层抽样，可以加上stratify_by_column="label"参数，替换为你自己的标签列名。

• 指定加载指定列：如果你的CSV有很多列，只需要保留需要的列可以加usecols参数：
  dataset = load_dataset("csv", data_files="./custom_data.csv", usecols=["text", "label"])
• 自定义分隔符：如果你的CSV用的不是逗号分隔，比如是制表符分隔的tsv文件，加sep="\t"参数
• 跳过损坏行：如果CSV里有格式错误的行，加on_bad_lines="skip"参数自动跳过

加载完数据集后可以直接配合transformers的工具做预处理，以文本分类的token化为例：

from transformers import AutoTokenizer

# 替换为你要使用的预训练模型名
tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")

def preprocess_func(examples):
    # 替换为你自己的文本列名
    return tokenizer(examples["text"], truncation=True, max_length=128)

# 批量处理整个数据集，batched参数开启批量处理速度更快
tokenized_dataset = dataset.map(preprocess_func, batched=True)

处理完成的数据集可以直接传给transformers.Trainer用于训练。

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