是时候剖析数据了

大多数时候，你使用的数据不会完美地为模型训练做好准备。在本节中，我们将探索 🤗 Datasets 提供的各种功能，以清理数据集。

剖析数据

与 Pandas 类似，🤗 Datasets 提供了多种函数来操作 `Dataset` 和 `DatasetDict` 对象。我们已经在第 3 章中遇到了 `Dataset.map()` 方法，在本节中，我们将探索我们可用的其他一些函数。

对于这个示例，我们将使用托管在加州大学欧文分校机器学习库上的药物评论数据集，其中包含患者对各种药物的评论，以及治疗的疾病和患者满意度的 10 星评级。

首先我们需要下载并解压数据，这可以通过 `wget` 和 `unzip` 命令完成。

!wget "https://archive.ics.uci.edu/ml/machine-learning-databases/00462/drugsCom_raw.zip"
!unzip drugsCom_raw.zip

由于 TSV 只是 CSV 的一种变体，使用制表符而不是逗号作为分隔符，我们可以通过使用 `csv` 加载脚本并在 `load_dataset()` 函数中指定 `delimiter` 参数来加载这些文件，如下所示：

from datasets import load_dataset

data_files = {"train": "drugsComTrain_raw.tsv", "test": "drugsComTest_raw.tsv"}
# \t is the tab character in Python
drug_dataset = load_dataset("csv", data_files=data_files, delimiter="\t")

在进行任何数据分析时，一个好的做法是抽取少量随机样本，以便快速了解正在处理的数据类型。在 🤗 Datasets 中，我们可以通过将 `Dataset.shuffle()` 和 `Dataset.select()` 函数链式调用来创建随机样本。

drug_sample = drug_dataset["train"].shuffle(seed=42).select(range(1000))
# Peek at the first few examples
drug_sample[:3]

{'Unnamed: 0': [87571, 178045, 80482],
 'drugName': ['Naproxen', 'Duloxetine', 'Mobic'],
 'condition': ['Gout, Acute', 'ibromyalgia', 'Inflammatory Conditions'],
 'review': ['"like the previous person mention, I&#039;m a strong believer of aleve, it works faster for my gout than the prescription meds I take. No more going to the doctor for refills.....Aleve works!"',
  '"I have taken Cymbalta for about a year and a half for fibromyalgia pain. It is great\r\nas a pain reducer and an anti-depressant, however, the side effects outweighed \r\nany benefit I got from it. I had trouble with restlessness, being tired constantly,\r\ndizziness, dry mouth, numbness and tingling in my feet, and horrible sweating. I am\r\nbeing weaned off of it now. Went from 60 mg to 30mg and now to 15 mg. I will be\r\noff completely in about a week. The fibro pain is coming back, but I would rather deal with it than the side effects."',
  '"I have been taking Mobic for over a year with no side effects other than an elevated blood pressure.  I had severe knee and ankle pain which completely went away after taking Mobic.  I attempted to stop the medication however pain returned after a few days."'],
 'rating': [9.0, 3.0, 10.0],
 'date': ['September 2, 2015', 'November 7, 2011', 'June 5, 2013'],
 'usefulCount': [36, 13, 128]}

请注意，我们已在 `Dataset.shuffle()` 中固定了种子以确保可重现性。`Dataset.select()` 需要一个可迭代的索引，因此我们传入了 `range(1000)` 来从打乱的数据集中获取前 1000 个示例。从这个样本中，我们已经可以看到数据集中存在一些奇怪之处：

“Unnamed: 0”列看起来很像每个患者的匿名 ID。
`condition` 列包含大写和小写标签的混合。
评论的长度各不相同，并且包含 Python 行分隔符（`\r\n`）以及 HTML 字符代码（如 `'`）。

让我们看看如何使用 🤗 Datasets 处理这些问题。为了验证 `Unnamed: 0` 列的患者 ID 假设，我们可以使用 `Dataset.unique()` 函数来验证 ID 数量是否与每个拆分中的行数匹配。

for split in drug_dataset.keys():
    assert len(drug_dataset[split]) == len(drug_dataset[split].unique("Unnamed: 0"))

这似乎证实了我们的假设，所以让我们通过将 `Unnamed: 0` 列重命名为更易于理解的名称来稍微清理一下数据集。我们可以使用 `DatasetDict.rename_column()` 函数一次性重命名两个拆分中的列。

drug_dataset = drug_dataset.rename_column(
    original_column_name="Unnamed: 0", new_column_name="patient_id"
)
drug_dataset

DatasetDict({
    train: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount'],
        num_rows: 161297
    })
    test: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount'],
        num_rows: 53766
    })
})

✏️ 试试看！ 使用 `Dataset.unique()` 函数找出训练集和测试集中独立药物和疾病的数量。

接下来，让我们使用 `Dataset.map()` 规范化所有 `condition` 标签。正如我们在第 3 章中对分词所做的那样，我们可以定义一个简单的函数，该函数可以应用于 `drug_dataset` 中每个分割的所有行。

def lowercase_condition(example):
    return {"condition": example["condition"].lower()}


drug_dataset.map(lowercase_condition)

AttributeError: 'NoneType' object has no attribute 'lower'

哦不，我们的映射函数出了问题！从错误中我们可以推断，`condition` 列中的某些条目是 `None`，它们不是字符串，因此无法转换为小写。让我们使用 `Dataset.filter()` 删除这些行，它的工作方式与 `Dataset.map()` 类似，并且需要一个接收数据集单个示例的函数。与其编写一个显式函数，例如：

def filter_nones(x):
    return x["condition"] is not None

然后运行 `drug_dataset.filter(filter_nones)`，我们可以使用lambda 函数在一行中完成此操作。在 Python 中，lambda 函数是可以不显式命名而定义的小函数。它们采用一般形式：

lambda <arguments> : <expression>

其中 `lambda` 是 Python 的一个特殊关键字，`<arguments>` 是一个逗号分隔的值列表/集合，用于定义函数的输入，`<expression>` 表示您希望执行的操作。例如，我们可以定义一个简单的 lambda 函数，用于计算数字的平方，如下所示：

lambda x : x * x

要将此函数应用于输入，我们需要将其与输入一起用括号括起来。

(lambda x: x * x)(3)

同样，我们可以通过逗号分隔参数来定义具有多个参数的 lambda 函数。例如，我们可以按如下方式计算三角形的面积：

(lambda base, height: 0.5 * base * height)(4, 8)

16.0

当你想要定义小型的、一次性使用的函数时，Lambda 函数非常方便（有关它们的更多信息，我们建议阅读 Andre Burgaud 撰写的优秀Real Python 教程）。在 🤗 Datasets 的上下文中，我们可以使用 Lambda 函数来定义简单的映射和过滤操作，所以让我们使用这个技巧来消除数据集中的 `None` 条目。

drug_dataset = drug_dataset.filter(lambda x: x["condition"] is not None)

删除了 `None` 条目后，我们可以规范化 `condition` 列。

drug_dataset = drug_dataset.map(lowercase_condition)
# Check that lowercasing worked
drug_dataset["train"]["condition"][:3]

['left ventricular dysfunction', 'adhd', 'birth control']

它奏效了！现在我们已经清理了标签，接下来让我们看看如何清理评论本身。

创建新列

每当你处理客户评论时，一个好的做法是检查每条评论中的单词数量。评论可能只有一个词，比如“太棒了！”，也可能是一篇包含数千词的完整论文，根据用例，你需要以不同的方式处理这些极端情况。为了计算每条评论中的单词数量，我们将使用基于按空格分割文本的粗略启发式方法。

让我们定义一个简单的函数来计算每条评论中的单词数。

def compute_review_length(example):
    return {"review_length": len(example["review"].split())}

与我们的 `lowercase_condition()` 函数不同，`compute_review_length()` 返回一个字典，其键与数据集中的列名不对应。在这种情况下，当 `compute_review_length()` 传递给 `Dataset.map()` 时，它将应用于数据集中的所有行，以创建一个新的 `review_length` 列。

drug_dataset = drug_dataset.map(compute_review_length)
# Inspect the first training example
drug_dataset["train"][0]

{'patient_id': 206461,
 'drugName': 'Valsartan',
 'condition': 'left ventricular dysfunction',
 'review': '"It has no side effect, I take it in combination of Bystolic 5 Mg and Fish Oil"',
 'rating': 9.0,
 'date': 'May 20, 2012',
 'usefulCount': 27,
 'review_length': 17}

正如预期的那样，我们可以看到一个 `review_length` 列已添加到我们的训练集中。我们可以使用 `Dataset.sort()` 对这个新列进行排序，以查看极端值是什么样子。

drug_dataset["train"].sort("review_length")[:3]

{'patient_id': [103488, 23627, 20558],
 'drugName': ['Loestrin 21 1 / 20', 'Chlorzoxazone', 'Nucynta'],
 'condition': ['birth control', 'muscle spasm', 'pain'],
 'review': ['"Excellent."', '"useless"', '"ok"'],
 'rating': [10.0, 1.0, 6.0],
 'date': ['November 4, 2008', 'March 24, 2017', 'August 20, 2016'],
 'usefulCount': [5, 2, 10],
 'review_length': [1, 1, 1]}

正如我们所怀疑的，有些评论只包含一个单词，这对于情感分析来说可能没问题，但如果我们要预测疾病，则不会提供信息。

🙋 向数据集添加新列的另一种方法是使用 `Dataset.add_column()` 函数。这允许您将列作为 Python 列表或 NumPy 数组提供，并且在 `Dataset.map()` 不太适合您的分析的情况下可能很有用。

让我们使用 `Dataset.filter()` 函数来删除包含少于 30 个单词的评论。类似于我们对 `condition` 列所做的，我们可以通过要求评论长度超过此阈值来过滤掉非常短的评论。

drug_dataset = drug_dataset.filter(lambda x: x["review_length"] > 30)
print(drug_dataset.num_rows)

{'train': 138514, 'test': 46108}

正如你所看到的，这已经从我们原来的训练集和测试集中删除了大约 15% 的评论。

✏️ 试一试！ 使用 `Dataset.sort()` 函数检查单词数量最多的评论。请参阅文档，了解您需要使用哪个参数来按降序排序评论的长度。

我们需要处理的最后一件事是评论中存在的 HTML 字符代码。我们可以使用 Python 的 `html` 模块来转义这些字符，如下所示：

import html

text = "I&#039;m a transformer called BERT"
html.unescape(text)

"I'm a transformer called BERT"

我们将使用 `Dataset.map()` 来解转义语料库中的所有 HTML 字符。

drug_dataset = drug_dataset.map(lambda x: {"review": html.unescape(x["review"])})

正如你所看到的，`Dataset.map()` 方法在处理数据方面非常有用——我们甚至还没有触及它能做的一切！

map() 方法的超能力

`Dataset.map()` 方法接受一个 `batched` 参数，如果设置为 `True`，则会一次性将一批示例发送到 map 函数（批大小可配置，但默认为 1,000）。例如，之前对所有 HTML 进行解转义的 map 函数运行时间稍长（您可以从进度条中读取运行时间）。我们可以通过使用列表推导式同时处理多个元素来加快速度。

当您指定 `batched=True` 时，函数会收到一个字典，其中包含数据集的字段，但每个值现在都是值列表，而不仅仅是单个值。`Dataset.map()` 的返回值也应该相同：一个字典，其中包含我们想要更新或添加到数据集的字段，以及一个值列表。例如，这是另一种解转义所有 HTML 字符的方法，但使用了 `batched=True`：

new_drug_dataset = drug_dataset.map(
    lambda x: {"review": [html.unescape(o) for o in x["review"]]}, batched=True
)

如果您在笔记本中运行此代码，您会发现此命令的执行速度比以前快得多。这并非因为我们的评论已经过 HTML 解转义——如果您重新执行上一节的指令（不带 `batched=True`），它将花费与以前相同的时间。这是因为列表推导式通常比在 `for` 循环中执行相同的代码更快，而且我们还可以通过同时访问大量元素而不是逐个访问来获得一些性能提升。

在 `Dataset.map()` 中使用 `batched=True` 对于释放我们将在第 6 章中遇到的“快速”分词器的速度至关重要，这些分词器可以快速地对大量文本进行分词。例如，要使用快速分词器对所有药物评论进行分词，我们可以使用如下函数：

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")


def tokenize_function(examples):
    return tokenizer(examples["review"], truncation=True)

正如你在第三章中看到的，我们可以向分词器传递一个或多个示例，因此我们可以有或没有 `batched=True` 来使用此函数。让我们借此机会比较不同选项的性能。在笔记本中，你可以通过在要测量的代码行前面添加 `%time` 来计时一行指令。

%time tokenized_dataset = drug_dataset.map(tokenize_function, batched=True)

您也可以通过在单元格开头放置 `%%time` 来计时整个单元格。在我们执行此操作的硬件上，此指令显示为 10.8 秒（这是“Wall time”后面写入的数字）。

✏️ 试一试！ 分别在带和不带 `batched=True` 的情况下执行相同的指令，然后尝试使用慢速分词器（在 `AutoTokenizer.from_pretrained()` 方法中添加 `use_fast=False`），这样你就可以在你的硬件上看到得到的数据。

以下是我们使用和不使用批处理，以及使用快速和慢速分词器获得的结果：

选项	快速分词器	慢速分词器
`batched=True`	10.8秒	4分41秒
`batched=False`	59.2秒	5分3秒

这意味着使用 `batched=True` 选项的快速分词器比没有批处理的慢速分词器快 30 倍——这真是令人惊叹！这是 `AutoTokenizer` 默认使用快速分词器（以及它们被称为“快速”的原因）的主要原因。它们能够实现如此快的速度提升，因为在幕后，分词代码是在 Rust 中执行的，Rust 是一种易于并行化代码执行的语言。

并行化也是快速分词器通过批处理实现近 6 倍速度提升的原因：您无法并行化单个分词操作，但当您想要同时对大量文本进行分词时，您可以将执行分拆到多个进程中，每个进程负责自己的文本。

`Dataset.map()` 自身也具有一些并行化功能。由于它们没有 Rust 的支持，因此无法让慢速分词器赶上快速分词器，但它们仍然很有用（特别是当您使用的分词器没有快速版本时）。要启用多处理，请使用 `num_proc` 参数，并在调用 `Dataset.map()` 时指定要使用的进程数。

slow_tokenizer = AutoTokenizer.from_pretrained("bert-base-cased", use_fast=False)


def slow_tokenize_function(examples):
    return slow_tokenizer(examples["review"], truncation=True)


tokenized_dataset = drug_dataset.map(slow_tokenize_function, batched=True, num_proc=8)

您可以稍微试验一下时间，以确定最佳的进程数；在我们的例子中，8 个似乎能产生最佳的速度增益。以下是我们使用和不使用多处理获得的数据：

选项	快速分词器	慢速分词器
`batched=True`	10.8秒	4分41秒
`batched=False`	59.2秒	5分3秒
`batched=True`，`num_proc=8`	6.52秒	41.3秒
`batched=False`，`num_proc=8`	9.49秒	45.2秒

这些对于慢速分词器来说是更合理的结果，但快速分词器的性能也得到了显著改善。然而，请注意，并非总是如此——对于 `num_proc` 的值不是 8 的情况，我们的测试表明使用 `batched=True` 但不使用该选项更快。通常，我们不建议对使用 `batched=True` 的快速分词器使用 Python 多处理。

使用 `num_proc` 来加速处理通常是一个很好的主意，只要您使用的函数本身不进行某种多处理。

所有这些功能都浓缩在一个方法中，这已经很了不起了，但还有更多！使用 `Dataset.map()` 和 `batched=True`，您可以更改数据集中的元素数量。这在许多情况下都非常有用，例如您想从一个示例创建多个训练特征，并且我们将在第 7 章中进行的几个 NLP 任务的预处理中需要这样做。

💡 在机器学习中，示例通常定义为我们输入到模型的特征集。在某些上下文中，这些特征将是 `Dataset` 中的列集，但在其他上下文中（例如这里和问答），可以从单个示例中提取多个特征，并且它们属于单个列。

让我们看看它是如何工作的！在这里，我们将对示例进行分词并将其截断为最大长度 128，但我们将要求分词器返回文本的所有块，而不仅仅是第一个。这可以通过 `return_overflowing_tokens=True` 来完成。

def tokenize_and_split(examples):
    return tokenizer(
        examples["review"],
        truncation=True,
        max_length=128,
        return_overflowing_tokens=True,
    )

在对整个数据集使用 `Dataset.map()` 之前，我们先在一个示例上进行测试。

result = tokenize_and_split(drug_dataset["train"][0])
[len(inp) for inp in result["input_ids"]]

[128, 49]

因此，训练集中的第一个示例变成了两个特征，因为它被分词后的标记数超过了我们指定的最大标记数：第一个长度为 128，第二个长度为 49。现在，让我们对数据集中的所有元素执行此操作！

tokenized_dataset = drug_dataset.map(tokenize_and_split, batched=True)

ArrowInvalid: Column 1 named condition expected length 1463 but got length 1000

哦不！这不起作用！为什么？查看错误消息会给我们一个线索：其中一列的长度不匹配，一列长度为 1,463，另一列长度为 1,000。如果您查看过 `Dataset.map()` 文档，您可能还记得它指的是传递给函数的样本数量；这里这 1,000 个示例产生了 1,463 个新特征，导致了形状错误。

问题在于我们试图混合两个不同大小的数据集：`drug_dataset` 列将有一定数量的示例（错误中的 1,000 个），但我们正在构建的 `tokenized_dataset` 将有更多（错误消息中的 1,463 个；它比 1,000 个多，因为我们通过使用 `return_overflowing_tokens=True` 将长评论分词成多个示例）。这对于 `Dataset` 不起作用，因此我们需要从旧数据集中删除列或使它们与新数据集中的大小相同。我们可以通过 `remove_columns` 参数完成前者。

tokenized_dataset = drug_dataset.map(
    tokenize_and_split, batched=True, remove_columns=drug_dataset["train"].column_names
)

现在这可以正常工作，没有错误。我们可以通过比较长度来检查我们的新数据集是否比原始数据集有更多元素。

len(tokenized_dataset["train"]), len(drug_dataset["train"])

(206772, 138514)

我们提到，我们还可以通过使旧列与新列的大小相同来处理长度不匹配的问题。为此，我们需要分词器在设置 `return_overflowing_tokens=True` 时返回的 `overflow_to_sample_mapping` 字段。它提供了从新特征索引到其原始样本索引的映射。使用这个，我们可以将原始数据集中存在的每个键与正确大小的值列表相关联，方法是将每个示例的值重复多次，因为它们生成了新的特征。

def tokenize_and_split(examples):
    result = tokenizer(
        examples["review"],
        truncation=True,
        max_length=128,
        return_overflowing_tokens=True,
    )
    # Extract mapping between new and old indices
    sample_map = result.pop("overflow_to_sample_mapping")
    for key, values in examples.items():
        result[key] = [values[i] for i in sample_map]
    return result

我们可以看到它在 `Dataset.map()` 中有效，而无需我们删除旧列。

tokenized_dataset = drug_dataset.map(tokenize_and_split, batched=True)
tokenized_dataset

DatasetDict({
    train: Dataset({
        features: ['attention_mask', 'condition', 'date', 'drugName', 'input_ids', 'patient_id', 'rating', 'review', 'review_length', 'token_type_ids', 'usefulCount'],
        num_rows: 206772
    })
    test: Dataset({
        features: ['attention_mask', 'condition', 'date', 'drugName', 'input_ids', 'patient_id', 'rating', 'review', 'review_length', 'token_type_ids', 'usefulCount'],
        num_rows: 68876
    })
})

我们获得了与之前相同的训练特征数量，但在这里我们保留了所有旧字段。如果您在应用模型后需要它们进行一些后处理，您可能希望使用这种方法。

你现在已经看到了如何使用 🤗 Datasets 以各种方式预处理数据集。尽管 🤗 Datasets 的处理功能将满足你大部分的模型训练需求，但有时你可能需要切换到 Pandas 以访问更强大的功能，例如 `DataFrame.groupby()` 或用于可视化的高级 API。幸运的是，🤗 Datasets 旨在与 Pandas、NumPy、PyTorch、TensorFlow 和 JAX 等库互操作。让我们看看这是如何工作的。

从 Dataset 到 DataFrame 再返回

为了实现各种第三方库之间的转换，🤗 Datasets 提供了 `Dataset.set_format()` 函数。此函数只更改数据集的输出格式，因此您可以轻松切换到其他格式，而不会影响底层数据格式，即 Apache Arrow。格式化是就地完成的。为了演示，让我们将数据集转换为 Pandas。

drug_dataset.set_format("pandas")

现在，当我们访问数据集的元素时，我们得到的是 `pandas.DataFrame` 而不是字典。

drug_dataset["train"][:3]

	患者ID	药物名称	状况	评论	评分	日期	有用计数	评论长度
0	95260	胍法辛	多动症	“我的儿子服用 Intuniv 已经四周了……”	8.0	2010年4月27日	192	141
1	92703	丽波乐	节育	“我以前服用过另一种口服避孕药，它有 21 天的周期，并且我非常满意——月经量非常少，最多 5 天，没有其他副作用……”	5.0	2009年12月14日	17	134
2	138000	奥索依芙拉	节育	“这是我第一次使用任何形式的避孕措施……”	8.0	2015年11月3日	10	89

让我们通过选择 `drug_dataset["train"]` 的所有元素来为整个训练集创建一个 `pandas.DataFrame`。

train_df = drug_dataset["train"][:]

🚨 在底层，`Dataset.set_format()` 会改变数据集 `__getitem__()` 魔术方法的返回格式。这意味着当我们想从处于 `"pandas"` 格式的 `Dataset` 创建像 `train_df` 这样的新对象时，我们需要对整个数据集进行切片以获取 `pandas.DataFrame`。您可以自己验证 `drug_dataset["train"]` 的类型是 `Dataset`，无论其输出格式如何。

从这里我们可以使用我们想要的所有 Pandas 功能。例如，我们可以进行花式链式调用来计算 `condition` 条目中的类别分布。

frequencies = (
    train_df["condition"]
    .value_counts()
    .to_frame()
    .reset_index()
    .rename(columns={"index": "condition", "count": "frequency"})
)
frequencies.head()

	状况	频率
0	节育	27655
1	抑郁症	8023
2	痤疮	5209
3	焦虑症	4991
4	疼痛	4744

一旦我们完成了 Pandas 分析，我们总是可以使用 `Dataset.from_pandas()` 函数创建新的 `Dataset` 对象，如下所示：

from datasets import Dataset

freq_dataset = Dataset.from_pandas(frequencies)
freq_dataset

Dataset({
    features: ['condition', 'frequency'],
    num_rows: 819
})

✏️ 试一试！ 计算每种药物的平均评分，并将结果存储在一个新的 `Dataset` 中。

到此为止，我们对 🤗 Datasets 中可用的各种预处理技术的介绍就结束了。为了总结本节，让我们创建一个验证集，以准备用于训练分类器的数据集。在此之前，我们将把 `drug_dataset` 的输出格式从 `"pandas"` 重置为 `"arrow"`。

drug_dataset.reset_format()

创建验证集

尽管我们有一个可用于评估的测试集，但在开发过程中保留测试集不变并创建一个单独的验证集是一个好习惯。一旦您对模型在验证集上的性能满意，您就可以对测试集进行最终的健全性检查。此过程有助于降低您过度拟合测试集并部署一个在实际数据上失败的模型的风险。

🤗 Datasets 提供了一个 `Dataset.train_test_split()` 函数，该函数基于 `scikit-learn` 中著名的功能。让我们使用它将训练集分成 `train` 和 `validation` 拆分（我们设置了 `seed` 参数以确保可重现性）。

drug_dataset_clean = drug_dataset["train"].train_test_split(train_size=0.8, seed=42)
# Rename the default "test" split to "validation"
drug_dataset_clean["validation"] = drug_dataset_clean.pop("test")
# Add the "test" set to our `DatasetDict`
drug_dataset_clean["test"] = drug_dataset["test"]
drug_dataset_clean

DatasetDict({
    train: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount', 'review_length', 'review_clean'],
        num_rows: 110811
    })
    validation: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount', 'review_length', 'review_clean'],
        num_rows: 27703
    })
    test: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount', 'review_length', 'review_clean'],
        num_rows: 46108
    })
})

太棒了，我们现在已经准备好一个数据集，可以用于训练一些模型了！在第 5 节中，我们将向您展示如何将数据集上传到 Hugging Face Hub，但现在，让我们通过查看几种将数据集保存到本地机器的方法来结束我们的分析。

保存数据集

尽管 🤗 Datasets 将缓存每个下载的数据集及其上执行的操作，但有时您会希望将数据集保存到磁盘（例如，以防缓存被删除）。如下表所示，🤗 Datasets 提供了三个主要函数来以不同格式保存您的数据集。

数据格式	函数
箭头	`Dataset.save_to_disk()`
CSV	`Dataset.to_csv()`
JSON	`Dataset.to_json()`

例如，让我们以 Arrow 格式保存清理后的数据集。

drug_dataset_clean.save_to_disk("drug-reviews")

这将创建一个具有以下结构的目录：

drug-reviews/
├── dataset_dict.json
├── test
│   ├── dataset.arrow
│   ├── dataset_info.json
│   └── state.json
├── train
│   ├── dataset.arrow
│   ├── dataset_info.json
│   ├── indices.arrow
│   └── state.json
└── validation
    ├── dataset.arrow
    ├── dataset_info.json
    ├── indices.arrow
    └── state.json

从中我们可以看到每个拆分都与它自己的 `dataset.arrow` 表以及 `dataset_info.json` 和 `state.json` 中的一些元数据相关联。您可以将 Arrow 格式视为一个花哨的列和行表，它经过优化，可用于构建处理和传输大型数据集的高性能应用程序。

数据集保存后，我们可以使用 `load_from_disk()` 函数加载它，如下所示：

from datasets import load_from_disk

drug_dataset_reloaded = load_from_disk("drug-reviews")
drug_dataset_reloaded

DatasetDict({
    train: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount', 'review_length'],
        num_rows: 110811
    })
    validation: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount', 'review_length'],
        num_rows: 27703
    })
    test: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount', 'review_length'],
        num_rows: 46108
    })
})

对于 CSV 和 JSON 格式，我们必须将每个拆分存储为单独的文件。一种方法是遍历 `DatasetDict` 对象中的键和值：

for split, dataset in drug_dataset_clean.items():
    dataset.to_json(f"drug-reviews-{split}.jsonl")

这会将每个分割保存为JSON Lines 格式，其中数据集中的每一行都存储为单独的 JSON 行。第一个示例看起来像这样：

!head -n 1 drug-reviews-train.jsonl

{"patient_id":141780,"drugName":"Escitalopram","condition":"depression","review":"\"I seemed to experience the regular side effects of LEXAPRO, insomnia, low sex drive, sleepiness during the day. I am taking it at night because my doctor said if it made me tired to take it at night. I assumed it would and started out taking it at night. Strange dreams, some pleasant. I was diagnosed with fibromyalgia. Seems to be helping with the pain. Have had anxiety and depression in my family, and have tried quite a few other medications that haven't worked. Only have been on it for two weeks but feel more positive in my mind, want to accomplish more in my life. Hopefully the side effects will dwindle away, worth it to stick with it from hearing others responses. Great medication.\"","rating":9.0,"date":"May 29, 2011","usefulCount":10,"review_length":125}

然后我们可以使用第 2 节中的技术加载 JSON 文件，如下所示：

data_files = {
    "train": "drug-reviews-train.jsonl",
    "validation": "drug-reviews-validation.jsonl",
    "test": "drug-reviews-test.jsonl",
}
drug_dataset_reloaded = load_dataset("json", data_files=data_files)

到此为止，我们对 🤗 Datasets 进行数据整理的探索就结束了！现在我们有了一个用于训练模型的清理数据集，以下是您可以尝试的一些想法：

使用第三章中的技术训练一个分类器，该分类器可以根据药物评论预测患者病情。
使用第一章中的 `summarization` 管道生成评论摘要。

接下来，我们将看看 🤗 Datasets 如何让您处理大型数据集而不会让您的笔记本电脑崩溃！

< > 在 GitHub 上更新