重要提示
您正在查看 NeMo 2.0 文档。此版本对 API 和新的库 NeMo Run 进行了重大更改。我们目前正在将 NeMo 1.0 的所有功能移植到 2.0。有关先前版本或 2.0 中尚未提供的功能的文档,请参阅 NeMo 24.07 文档。
数据集#
DocumentDataset#
- class nemo_curator.datasets.DocumentDataset(dataset_df: dask.dataframe.DataFrame)#
文档和文档元数据的集合。在内部,它可能分布在多个节点上,并且可能在 GPU 上。
- classmethod from_pandas(
- 数据,
- npartitions: int | None = 1,
- chunksize: int | None = None,
- sort: bool | None = True,
- name: str | None = None,
从 pandas 数据帧创建文档数据集。有关参数的更多信息,请参阅 Dask 的 from_pandas 文档 https://docs.dask.org/en/stable/generated/dask.dataframe.from_pandas.html
- 参数:
data – pandas 数据帧
- 返回值:
具有 pandas 后端(在 CPU 上)的文档数据集。
- classmethod read_json(
- input_files: str | List[str],
- backend: Literal['pandas', 'cudf'] = 'pandas',
- files_per_partition: int | None = None,
- blocksize: str | None = '1gb',
- add_filename: bool | str = False,
- input_meta: str | dict | None = None,
- columns: List[str] | None = None,
- **kwargs,
读取 JSONL 或 JSONL 文件。
- 参数:
input_files – 输入文件的路径。
backend – 用于读取数据的后端。
files_per_partition – 每个分区读取的文件数。
add_filename – 是否向 DataFrame 添加文件名列。如果为 True,则向 DataFrame 添加一个名为 file_name 的新列。如果为 str,则设置新的列名。默认为 False。
input_meta – 一个字典或格式化为字典的字符串,概述了 JSONL 输入文件中的字段名称及其各自的数据类型。
columns – 如果不为 None,则仅从文件中读取这些列。
- classmethod read_parquet(
- input_files: str | List[str],
- backend: Literal['pandas', 'cudf'] = 'pandas',
- files_per_partition: int | None = None,
- blocksize: str | None = '1gb',
- add_filename: bool | str = False,
- columns: List[str] | None = None,
- **kwargs,
读取 Parquet 文件。
- 参数:
input_files – 输入文件的路径。
backend – 用于读取数据的后端。
files_per_partition – 每个分区读取的文件数。
add_filename – 是否向 DataFrame 添加文件名列。如果为 True,则向 DataFrame 添加一个名为 file_name 的新列。如果为 str,则设置新的列名。默认为 False。
columns – 如果不为 None,则仅从文件中读取这些列。在为 Parquet 文件指定列时,性能会显着提高。
- classmethod read_pickle(
- input_files: str | List[str],
- backend: Literal['pandas', 'cudf'] = 'pandas',
- columns: List[str] | None = None,
- **kwargs,
读取 Pickle 文件。
- 参数:
input_files – 输入文件的路径。
backend – 用于读取数据的后端。
files_per_partition – 每个分区读取的文件数。
add_filename – 是否向 DataFrame 添加文件名列。如果为 True,则向 DataFrame 添加一个名为 file_name 的新列。如果为 str,则设置新的列名。默认为 False。
columns – 如果不为 None,则仅从文件中读取这些列。
- to_json(
- output_path: str,
- write_to_filename: bool | str = False,
- keep_filename_column: bool = False,
- partition_on: str | None = None,
以 JSONL 格式将数据集写入到指定路径。
如果 write_to_filename 为 True,则 DataFrame 应该有一列指定每个文档的文件名。默认情况下,此列可以命名为 file_name,或者如果 write_to_filename 是字符串,则可以使用自定义名称。
- 参数:
output_path (str) – 将数据集写入到的目录或文件路径。
write_to_filename (Union[bool, str]) – 确定如何处理文件名。- 如果为 True,则使用 DataFrame 中的 file_name 列来确定文件名。- 如果为字符串,则使用该字符串作为文件名的列名。- 如果为 False,则将所有数据写入到指定的 output_path。
keep_filename_column (bool) – 如果为 True,则在输出中保留文件名列。如果为 False,则从输出中删除文件名列。
partition_on (Optional[str]) – 用于对数据进行分区的列名。如果指定,则根据此列中的唯一值对数据进行分区,每个分区写入到单独的目录。
有关更多详细信息,请参阅 nemo_curator.utils.distributed_utils 中的 write_to_disk 函数。
- to_pandas()#
从 DocumentDataset 创建 pandas 数据帧
- 返回值:
pandas 数据帧(在 CPU 上)
- to_parquet(
- output_path: str,
- write_to_filename: bool | str = False,
- keep_filename_column: bool = False,
- partition_on: str | None = None,
以 Parquet 格式将数据集写入到指定路径。
如果 write_to_filename 为 True,则 DataFrame 应该有一列指定每个文档的文件名。默认情况下,此列可以命名为 file_name,或者如果 write_to_filename 是字符串,则可以使用自定义名称。
- 参数:
output_path (str) – 将数据集写入到的目录或文件路径。
write_to_filename (Union[bool, str]) – 确定如何处理文件名。- 如果为 True,则使用 DataFrame 中的 file_name 列来确定文件名。- 如果为字符串,则使用该字符串作为文件名的列名。- 如果为 False,则将所有数据写入到指定的 output_path。
keep_filename_column (bool) – 如果为 True,则在输出中保留文件名列。如果为 False,则从输出中删除文件名列。
partition_on (Optional[str]) – 用于对数据进行分区的列名。如果指定,则根据此列中的唯一值对数据进行分区,每个分区写入到单独的目录。
有关更多详细信息,请参阅 nemo_curator.utils.distributed_utils 中的 write_to_disk 函数。
- class nemo_curator.datasets.ParallelDataset(dataset_df: dask.dataframe.DataFrame)#
标准 DocumentDataset 的扩展,具有加载简单双语文本的特殊方法。
对于具有更复杂元数据的数据,请将您的数据转换为 jsonl/parquet/pickle 格式,并使用 DocumentDataset 中定义的接口。
- classmethod read_simple_bitext(
- src_input_files: str | List[str],
- tgt_input_files: str | List[str],
- src_lang: str,
- tgt_lang: str,
- backend: str = 'pandas',
- add_filename: bool | str = False,
- npartitions: int = 16,
有关 “simple_bitext” 的含义以及其他参数的用法,请参阅 read_single_simple_bitext_file_pair 文档字符串。
- 参数:
src_input_files (Union[str, List[str]]) – 一个或多个输入文件,以源语言
tgt_input_files (Union[str, List[str]]) – 一个或多个输入文件,以目标语言
- Raises:
TypeError – 如果 src_input_files 和 tgt_input_files 的类型不一致。
- 返回值:
一个 ParallelDataset 对象,其中 self.df 包含摄取的简单双语文本。
- Return type:
- static read_single_simple_bitext_file_pair(
- input_file_pair: Tuple[str],
- src_lang: str,
- tgt_lang: str,
- doc_id: str = None,
- backend: str = 'cudf',
- add_filename: bool | str = False,
此函数将一对 “简单双语文本” 文件读取到 pandas DataFrame 中。简单双语文本是机器翻译中常用的数据格式。它由两个纯文本文件组成,文件行数相同,每行对是彼此的翻译。例如
data.de
` Wir besitzen keine Reisetaschen aus Leder. Die Firma produziert Computer für den deutschen Markt. ... `
data.en
` We don't own duffel bags made of leather. The company produces computers for the German market. ... `
为简单起见,我们还假设两个文本文件的名称具有相同的前缀,除了末尾的文件扩展名不同语言代码。
- 参数:
input_file_pair (Tuple[str]) – 指向输入文件的一对文件路径
src_lang (str) – 源语言,采用 ISO-639-1(两个字符)格式(例如 ‘en’)
tgt_lang (str) – 目标语言,采用 ISO-639-1(两个字符)格式(例如 ‘en’)
doc_id (str, optional) – 要分配给文件中每个段的字符串文档 ID。默认为 None。
backend (str, optional) – 数据帧的后端。默认为 “cudf”。
add_filename (Union[bool, str]) – 是否向 DataFrame 添加文件名列。如果为 True,则向 DataFrame 添加一个名为 file_name 的新列。如果为 str,则设置新的列名。默认为 False。
- 返回值:
Union[dd.DataFrame, dask_cudf.DataFrame]
- to_bitext(output_file_dir, write_to_filename=False)#
有关参数用法,请参阅 nemo_curator.utils.distributed_utils.write_to_disk 文档字符串。
ImageTextPairDataset#
- class nemo_curator.datasets.ImageTextPairDataset(
- path: str,
- metadata: dask.dataframe.DataFrame,
- tar_files: List[str],
- id_col: str,
以类似 WebDataset 的格式存储在磁盘或云存储上的图像文本对的集合。
确切的格式假定一个目录,其中包含分片的 .tar、.parquet 和(可选) .idx 文件。每个 tar 文件都应具有唯一的整数 ID 作为其名称(00000.tar、00001.tar、00002.tar 等)。tar 文件应包含 .jpg 文件中的图像、.txt 文件中的文本标题和 .json 文件中的元数据。数据集的每个记录都由唯一的 ID 标识,该 ID 是分片 ID 以及记录在分片中的偏移量的组合。例如,第 43 个分片的第 32 个记录将在 00042.tar 中,并具有图像 000420031.jpg、标题 000420031.txt 和元数据 000420031.json(假设从零开始索引)。
除了 tar 文件集合之外,ImageTextPairDataset 还期望根目录中存在 .parquet 文件,这些文件遵循与分片相同的命名约定(00042.tar -> 00042.parquet)。每个 Parquet 文件都应包含每个记录的元数据的聚合表格形式,其中 Parquet 文件中的每一行对应于该分片中的一个记录。Parquet 文件和 JSON 文件中的元数据都必须包含一个唯一的 ID 列,该列与其记录 ID 相同(在我们的示例中为 000420031)。
索引文件也可能位于目录中,以加快使用 DALI 的数据加载速度。索引文件必须由 DALI 的 wds2idx 工具生成。有关更多信息,请参阅 http://docs.nvda.net.cn/deeplearning/dali/user-guide/docs/examples/general/data_loading/dataloading_webdataset.html#Creating-an-index。每个索引文件都必须遵循与 tar 文件相同的命名约定(00042.tar -> 00042.idx)。
- classmethod from_webdataset(path: str, id_col: str)#
从 WebDataset 加载 ImageTextPairDataset
- 参数:
path (str) – WebDataset 样式的格式在磁盘或云存储上的路径。
id_col (str) – 存储每个记录的唯一标识符的列。
- save_metadata(
- path: str | None = None,
- columns: List[str] | None = None,
将数据集的元数据保存到指定路径作为 Parquet 文件集合。
- 参数:
path (Optional[str]) – 要将元数据保存到的路径。如果为 None,则写入到原始路径。
columns (Optional[List[str]]) – 如果指定,则仅保存列的子集。
- to_webdataset(
- path: str,
- filter_column: str,
- samples_per_shard: int = 10000,
- max_shards: int = 5,
- old_id_col: str | None = None,
将数据集保存为带有 Parquet 文件的 WebDataset 格式。 将重新分片 tar 文件,使其每个分片包含指定数量的样本。 ImageTextPairDataset.id_col 中的 ID 值将被新的 ID 覆盖。
- 参数:
path (str) – 数据集应写入的输出路径。
filter_column (str) – 一个布尔值列。 此列中值为 True 的所有样本都将包含在输出中。 否则,将省略该样本。
samples_per_shard (int) – 每个 tar 文件中包含的样本数。
max_shards (int) – 将从数据集中创建的最大分片数量级。 将用于确定分片/样本 ID 中前导零的数量。
old_id_col (Optional[str]) – 如果指定,将在给定列中保留先前的 ID 值。