Python Datasets的FolderBasedBuilder 类 处理jsonl的底层代码
flyfish
src\datasets\packaged_modules\folder_based_builder\folder_based_builder.py
主要功能
FolderBasedBuilder 类提供了一种通用的方式来从文件夹中加载数据，并可以处理数据文件的元数据。它支持多种文件格式（如 CSV、JSONL 和 Parquet）的元数据文件，同时可以根据文件夹结构推断数据的标签。
folder_based_builder.py 中的代码定义了一个名为 FolderBasedBuilder 的基类，它主要用于构建通用的数据加载器，特别是针对视觉和图像数据。
具体功能点
1. 类属性定义
BASE_FEATURE: 用于解码数据的特征对象，例如datasets.Image或datasets.Audio。BASE_COLUMN_NAME: 基础特征的键名，例如"image"或"audio"。BUILDER_CONFIG_CLASS: 继承自folder_based_builder.FolderBasedBuilderConfig的构建器配置类。EXTENSIONS: 允许的文件扩展名列表，只有具有这些扩展名的文件和元数据文件会被包含在数据集中。METADATA_FILENAMES: 元数据文件的文件名列表，包括"metadata.csv"、"metadata.jsonl"和"metadata.parquet"。
2. 数据加载和分割
_info(self): 返回数据集的信息，包括特征定义。_split_generators(self, dl_manager): 根据配置的文件路径下载和提取数据文件，并根据需要分析数据以推断标签和查找元数据文件。它会将数据划分为不同的数据集分割（如训练集、验证集和测试集）。_split_files_and_archives(self, data_files): 将数据文件分为普通文件和存档文件。
3. 元数据读取
_read_metadata(self, metadata_file: str, metadata_ext: str = ""): 根据元数据文件的扩展名（CSV、JSONL 或 Parquet），使用不同的方法读取元数据文件，并以pyarrow.Table的形式返回数据块。
4. 数据生成
_generate_examples(self, files, metadata_files, add_metadata, add_labels): 生成数据集中的样本。如果有元数据文件，则使用元数据来生成样本；否则，根据文件夹结构生成样本，并可以选择添加标签。
代码示例中的使用场景
在代码示例中，FolderBasedBuilder 类被用于构建一个虚拟的数据集加载器 DummyFolderBasedBuilder，该加载器用于加载 .txt 文件，并可以处理相关的元数据文件。
class DummyFolderBasedBuilder(FolderBasedBuilder):
BASE_FEATURE = DummyFeature
BASE_COLUMN_NAME = "base"
BUILDER_CONFIG_CLASS = FolderBasedBuilderConfig
EXTENSIONS = [".txt"]
import collections
import io
import itertools
import os
from dataclasses import dataclass
from typing import Any, Callable, Iterator, Optional, Union
import pandas as pd
import pyarrow as pa
import pyarrow.dataset as ds
import pyarrow.json as paj
import pyarrow.parquet as pq
import datasets
from datasets import config
from datasets.features.features import FeatureType, _visit, _visit_with_path, _VisitPath, require_storage_cast
from datasets.utils.file_utils import readline
# 获取日志记录器，用于记录程序运行时的信息
logger = datasets.utils.logging.get_logger(__name__)
# 定义一个函数，用于计算路径中的分隔符数量，这里使用正斜杠作为分隔符进行统计
def count_path_segments(path):
return path.replace("\\", "/").count("/")
# 定义一个数据类，继承自 datasets.BuilderConfig，用于配置 FolderBasedBuilder 的参数
@dataclass
class FolderBasedBuilderConfig(datasets.BuilderConfig):
"""BuilderConfig for AutoFolder."""
# 可选的特征描述，用于指定数据集的特征
features: Optional[datasets.Features] = None
# 是否丢弃标签的标志，默认为 None
drop_labels: bool = None
# 是否丢弃元数据的标志，默认为 None
drop_metadata: bool = None
# 可选的过滤器，用于过滤数据
filters: Optional[Union[ds.Expression, list[tuple], list[list[tuple]]]] = None
def __post_init__(self):
# 调用父类的 __post_init__ 方法
super().__post_init__()
# 定义一个基类，继承自 datasets.GeneratorBasedBuilder，用于通用的数据加载器，主要处理视觉和图像数据
class FolderBasedBuilder(datasets.GeneratorBasedBuilder):
"""
Base class for generic data loaders for vision and image data.
Abstract class attributes to be overridden by a child class:
BASE_FEATURE: feature object to decode data (i.e. datasets.Image, datasets.Audio, ...)
BASE_COLUMN_NAME: string key name of a base feature (i.e. "image", "audio", ...)
BUILDER_CONFIG_CLASS: builder config inherited from `folder_based_builder.FolderBasedBuilderConfig`
EXTENSIONS: list of allowed extensions (only files with these extensions and METADATA_FILENAME files
will be included in a dataset)
"""
# 抽象属性，用于指定解码数据的特征对象，需要子类进行重写
BASE_FEATURE: type[FeatureType]
# 抽象属性，用于指定基础特征的字符串键名，需要子类进行重写
BASE_COLUMN_NAME: str
# 抽象属性，用于指定继承自 FolderBasedBuilderConfig 的构建器配置类，需要子类进行重写
BUILDER_CONFIG_CLASS: FolderBasedBuilderConfig
# 抽象属性，用于指定允许的文件扩展名列表，需要子类进行重写
EXTENSIONS: list[str]
# 元数据文件名列表，包含支持的元数据文件类型
METADATA_FILENAMES: list[str] = ["metadata.csv", "metadata.jsonl", "metadata.parquet"]
# 返回数据集的信息，主要包含特征描述
def _info(self):
return datasets.DatasetInfo(features=self.config.features)
# 分割数据集生成器，将数据分割成不同的数据集划分（如训练集、验证集、测试集）
def _split_generators(self, dl_manager):
# 检查是否指定了数据文件，如果没有则抛出异常
if not self.config.data_files:
raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
# 设置下载管理器的配置，开启实时解压
dl_manager.download_config.extract_on_the_fly = True
# 判断是否需要进行早期分析，即是否需要推断类标签或查找元数据文件
do_analyze = not self.config.drop_labels or not self.config.drop_metadata
# 用于存储类标签的集合
labels, path_depths = set(), set()
# 用于存储每个数据集划分的元数据文件的字典
metadata_files = collections.defaultdict(set)
# 定义一个内部函数，用于分析文件或存档，查找标签和元数据文件
def analyze(files_or_archives, downloaded_files_or_dirs, split):
# 如果下载的文件或目录列表为空，则直接返回
if len(downloaded_files_or_dirs) == 0:
return
# 检查第一个样本是文件还是目录，以决定如何迭代
if os.path.isfile(downloaded_files_or_dirs[0]):
original_files, downloaded_files = files_or_archives, downloaded_files_or_dirs
for original_file, downloaded_file in zip(original_files, downloaded_files):
original_file, downloaded_file = str(original_file), str(downloaded_file)
# 获取原始文件的扩展名
_, original_file_ext = os.path.splitext(original_file)
# 如果文件扩展名在允许的扩展名列表中
if original_file_ext.lower() in self.EXTENSIONS:
# 如果不丢弃标签，则将标签添加到集合中，并记录路径深度
if not self.config.drop_labels:
labels.add(os.path.basename(os.path.dirname(original_file)))
path_depths.add(count_path_segments(original_file))
# 如果文件名是元数据文件名之一，则将其添加到元数据文件集合中
elif os.path.basename(original_file) in self.METADATA_FILENAMES:
metadata_files[split].add((original_file, downloaded_file))
else:
# 记录被忽略的文件信息
original_file_name = os.path.basename(original_file)
logger.debug(
f"The file '{original_file_name}' was ignored: it is not a {self.BASE_COLUMN_NAME}, and is not {self.METADATA_FILENAMES} either."
)
else:
archives, downloaded_dirs = files_or_archives, downloaded_files_or_dirs
for archive, downloaded_dir in zip(archives, downloaded_dirs):
archive, downloaded_dir = str(archive), str(downloaded_dir)
for downloaded_dir_file in dl_manager.iter_files(downloaded_dir):
# 获取下载文件的扩展名
_, downloaded_dir_file_ext = os.path.splitext(downloaded_dir_file)
# 如果文件扩展名在允许的扩展名列表中
if downloaded_dir_file_ext in self.EXTENSIONS:
# 如果不丢弃标签，则将标签添加到集合中，并记录路径深度
if not self.config.drop_labels:
labels.add(os.path.basename(os.path.dirname(downloaded_dir_file)))
path_depths.add(count_path_segments(downloaded_dir_file))
# 如果文件名是元数据文件名之一，则将其添加到元数据文件集合中
elif os.path.basename(downloaded_dir_file) in self.METADATA_FILENAMES:
metadata_files[split].add((None, downloaded_dir_file))
else:
# 记录被忽略的文件信息
archive_file_name = os.path.basename(archive)
original_file_name = os.path.basename(downloaded_dir_file)
logger.debug(
f"The file '{original_file_name}' from the archive '{archive_file_name}' was ignored: it is not a {self.BASE_COLUMN_NAME}, and is not {self.METADATA_FILENAMES} either."
)
# 获取配置中的数据文件
data_files = self.config.data_files
# 用于存储数据集划分生成器的列表
splits = []
for split_name, files in data_files.items():
# 如果文件是字符串，则将其转换为列表
if isinstance(files, str):
files = [files]
# 将数据文件分为普通文件和存档文件
files, archives = self._split_files_and_archives(files)
# 下载普通文件
downloaded_files = dl_manager.download(files)
# 下载并解压存档文件
downloaded_dirs = dl_manager.download_and_extract(archives)
# 如果需要进行早期分析
if do_analyze: # drop_metadata is None or False, drop_labels is None or False
# 记录正在搜索标签和/或元数据文件的信息
logger.info(f"Searching for labels and/or metadata files in {split_name} data files...")
# 分析普通文件
analyze(files, downloaded_files, split_name)
# 分析存档文件
analyze(archives, downloaded_dirs, split_name)
# 如果找到了元数据文件
if metadata_files:
# 如果不丢弃元数据，则添加元数据
add_metadata = not self.config.drop_metadata
# 如果找到了元数据文件，则不添加标签
add_labels = False
else:
# 如果没有找到元数据文件，则不添加元数据
add_metadata = False
# 如果没有找到元数据文件且不丢弃标签，并且满足一定条件，则添加标签
add_labels = (
(len(labels) > 1 and len(path_depths) == 1)
if self.config.drop_labels is None
else not self.config.drop_labels
)
# 如果添加标签，则记录信息
if add_labels:
logger.info("Adding the labels inferred from data directories to the dataset's features...")
# 如果添加元数据，则记录信息
if add_metadata:
logger.info("Adding metadata to the dataset...")
else:
# 如果不需要进行早期分析，则不添加标签和元数据
add_labels, add_metadata, metadata_files = False, False, {}
# 将数据集划分生成器添加到列表中
splits.append(
datasets.SplitGenerator(
name=split_name,
gen_kwargs={
"files": tuple(zip(files, downloaded_files))
+ tuple((None, dl_manager.iter_files(downloaded_dir)) for downloaded_dir in downloaded_dirs),
"metadata_files": metadata_files.get(split_name, []),
"add_labels": add_labels,
"add_metadata": add_metadata,
},
)
)
# 如果添加元数据
if add_metadata:
# 用于存储每个元数据文件的特征信息的列表
features_per_metadata_file: list[tuple[str, datasets.Features]] = []
# 检查所有元数据文件是否具有相同的扩展名
metadata_ext = {
os.path.splitext(original_metadata_file or downloaded_metadata_file)[-1]
for original_metadata_file, downloaded_metadata_file in itertools.chain.from_iterable(
metadata_files.values()
)
}
# 如果元数据文件有不同的扩展名，则抛出异常
if len(metadata_ext) > 1:
raise ValueError(f"Found metadata files with different extensions: {list(metadata_ext)}")
# 获取元数据文件的扩展名
metadata_ext = metadata_ext.pop()
# 遍历每个数据集划分的元数据文件
for split_metadata_files in metadata_files.values():
pa_metadata_table = None
for _, downloaded_metadata_file in split_metadata_files:
# 读取元数据文件的第一行
for pa_metadata_table in self._read_metadata(downloaded_metadata_file, metadata_ext=metadata_ext):
break # just fetch the first rows
if pa_metadata_table is not None:
# 将元数据文件的特征信息添加到列表中
features_per_metadata_file.append(
(downloaded_metadata_file, datasets.Features.from_arrow_schema(pa_metadata_table.schema))
)
break # no need to fetch all the files
# 检查所有元数据文件的特征是否相同
for downloaded_metadata_file, metadata_features in features_per_metadata_file:
if metadata_features != features_per_metadata_file[0][1]:
raise ValueError(
f"Metadata files {downloaded_metadata_file} and {features_per_metadata_file[0][0]} have different features: {features_per_metadata_file[0]} != {metadata_features}"
)
# 获取元数据文件的特征信息
metadata_features = features_per_metadata_file[0][1]
# 用于标记是否找到特定特征的标志
feature_not_found = True
# 定义一个内部函数，用于设置特征
def _set_feature(feature):
nonlocal feature_not_found
if isinstance(feature, dict):
out = type(feature)()
for key in feature:
# 如果键是 "file_name" 或以 "_file_name" 结尾，并且值的类型是字符串
if (key == "file_name" or key.endswith("_file_name")) and feature[key] == datasets.Value(
"string"
):
# 处理键名
key = key[: -len("_file_name")] or self.BASE_COLUMN_NAME
# 设置特征
out[key] = self.BASE_FEATURE()
feature_not_found = False
# 如果键是 "file_names" 或以 "_file_names" 结尾，并且值的类型是字符串序列
elif (key == "file_names" or key.endswith("_file_names")) and feature[
key
] == datasets.Sequence(datasets.Value("string")):
# 处理键名
key = key[: -len("_file_names")] or (self.BASE_COLUMN_NAME + "s")
# 设置特征
out[key] = datasets.Sequence(self.BASE_FEATURE())
feature_not_found = False
# 如果键是 "file_names" 或以 "_file_names" 结尾，并且值的类型是字符串列表
elif (key == "file_names" or key.endswith("_file_names")) and feature[key] == [
datasets.Value("string")
]:
# 处理键名
key = key[: -len("_file_names")] or (self.BASE_COLUMN_NAME + "s")
# 设置特征
out[key] = [self.BASE_FEATURE()]
feature_not_found = False
else:
# 其他情况，直接复制特征
out[key] = feature[key]
return out
return feature
# 对元数据特征进行处理
metadata_features = _visit(metadata_features, _set_feature)
# 如果没有找到特定特征，则抛出异常
if feature_not_found:
raise ValueError(
"`file_name` or `*_file_name` must be present as dictionary key (with type string) in metadata files"
)
else:
# 如果不添加元数据，则元数据特征为 None
metadata_features = None
# 如果配置中没有指定特征
if self.config.features is None:
if add_metadata:
# 如果添加元数据，则使用元数据特征
self.info.features = metadata_features
elif add_labels:
# 如果添加标签，则构建包含基础特征和标签的特征信息
self.info.features = datasets.Features(
{
self.BASE_COLUMN_NAME: self.BASE_FEATURE(),
"label": datasets.ClassLabel(names=sorted(labels)),
}
)
else:
# 否则，只包含基础特征
self.info.features = datasets.Features({self.BASE_COLUMN_NAME: self.BASE_FEATURE()})
# 返回数据集划分生成器列表
return splits
# 将数据文件分为普通文件和存档文件
def _split_files_and_archives(self, data_files):
files, archives = [], []
for data_file in data_files:
# 获取数据文件的扩展名
_, data_file_ext = os.path.splitext(data_file)
# 如果文件扩展名在允许的扩展名列表中，或者文件名是元数据文件名之一，则添加到普通文件列表中
if data_file_ext.lower() in self.EXTENSIONS or os.path.basename(data_file) in self.METADATA_FILENAMES:
files.append(data_file)
else:
# 否则，添加到存档文件列表中
archives.append(data_file)
return files, archives
# 读取元数据文件，并以 PyArrow 表的形式返回
def _read_metadata(self, metadata_file: str, metadata_ext: str = "") -> Iterator[pa.Table]:
"""using the same logic as the Csv, Json and Parquet dataset builders to stream the data"""
# 如果配置中指定了过滤器，则将其转换为 PyArrow 表达式
if self.config.filters is not None:
filter_expr = (
pq.filters_to_expression(self.config.filters)
if isinstance(self.config.filters, list)
else self.config.filters
)
else:
# 否则，过滤器表达式为 None
filter_expr = None
# 如果元数据文件是 CSV 格式
if metadata_ext == ".csv":
# 每次读取 10000 行
chunksize = 10_000
# 获取配置中的特征模式
schema = self.config.features.arrow_schema if self.config.features else None
# 用于指定读取 CSV 文件时的数据类型
dtype = (
{
name: dtype.to_pandas_dtype() if not require_storage_cast(feature) else object
for name, dtype, feature in zip(schema.names, schema.types, self.config.features.values())
}
if schema is not None
else None
)
# 使用 Pandas 读取 CSV 文件，以迭代器的形式返回
csv_file_reader = pd.read_csv(metadata_file, iterator=True, dtype=dtype, chunksize=chunksize)
for df in csv_file_reader:
# 将 Pandas 数据框转换为 PyArrow 表
pa_table = pa.Table.from_pandas(df)
# 如果有过滤器，则对表进行过滤
if self.config.filters is not None:
pa_table = pa_table.filter(filter_expr)
# 如果表不为空，则返回该表
if len(pa_table) > 0:
yield pa_table
# 如果元数据文件是 JSONL 格式
elif metadata_ext == ".jsonl":
with open(metadata_file, "rb") as f:
# 每次读取 10MB 的数据
chunksize: int = 10 << 20
# 用于多线程处理的块大小，最小为 16kB
block_size = max(chunksize // 32, 16 << 10)
while True:
# 读取指定大小的块
batch = f.read(chunksize)
if not batch:
break
# 读取当前行的剩余部分
try:
batch += f.readline()
except (AttributeError, io.UnsupportedOperation):
batch += readline(f)
while True:
try:
# 使用 PyArrow 读取 JSON 数据
pa_table = paj.read_json(
io.BytesIO(batch), read_options=paj.ReadOptions(block_size=block_size)
)
break
except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
# 如果出现特定错误，并且块大小已经超过了批量大小，则抛出异常
if (
isinstance(e, pa.ArrowInvalid)
and "straddling" not in str(e)
or block_size > len(batch)
):
raise
else:
# 否则，增加块大小并重试
logger.debug(
f"Batch of {len(batch)} bytes couldn't be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}."
)
block_size *= 2
# 如果有过滤器，则对表进行过滤
if self.config.filters is not None:
pa_table = pa_table.filter(filter_expr)
# 如果表不为空，则返回该表
if len(pa_table) > 0:
yield pa_table
else:
with open(metadata_file, "rb") as f:
# 创建 Parquet 文件片段
parquet_fragment = ds.ParquetFileFormat().make_fragment(f)
if parquet_fragment.row_groups:
# 如果有行组，则使用第一个行组的行数作为批量大小
batch_size = parquet_fragment.row_groups[0].num_rows
else:
# 否则，使用默认的最大批量大小
batch_size = config.DEFAULT_MAX_BATCH_SIZE
# 以批量的形式读取 Parquet 文件
for record_batch in parquet_fragment.to_batches(
batch_size=batch_size,
filter=filter_expr,
batch_readahead=0,
fragment_readahead=0,
):
# 将记录批量转换为 PyArrow 表并返回
yield pa.Table.from_batches([record_batch])
# 生成数据集的样本
def _generate_examples(self, files, metadata_files, add_metadata, add_labels):
# 样本索引
sample_idx = 0
# 如果添加元数据
if add_metadata:
# 用于存储特征路径的列表
feature_paths = []
# 定义一个内部函数，用于查找特征路径
def find_feature_path(feature, feature_path):
nonlocal feature_paths
if feature_path and isinstance(feature, self.BASE_FEATURE):
feature_paths.append(feature_path)
# 遍历特征信息，查找特征路径
_visit_with_path(self.info.features, find_feature_path)
for original_metadata_file, downloaded_metadata_file in metadata_files:
# 获取元数据文件的扩展名
metadata_ext = os.path.splitext(original_metadata_file or downloaded_metadata_file)[-1]
# 获取下载的元数据文件所在的目录
downloaded_metadata_dir = os.path.dirname(downloaded_metadata_file)
# 定义一个内部函数，用于设置特征值
def set_feature(item, feature_path: _VisitPath):
if len(feature_path) == 2 and isinstance(feature_path[0], str) and feature_path[1] == 0:
item[feature_path[0]] = item.pop("file_names", None) or item.pop(
feature_path[0] + "_file_names", None
)
elif len(feature_path) == 1 and isinstance(feature_path[0], str):
item[feature_path[0]] = item.pop("file_name", None) or item.pop(
feature_path[0] + "_file_name", None
)
elif len(feature_path) == 0:
# 处理文件相对路径
file_relpath = os.path.normpath(item).replace("\\", "/")
item = os.path.join(downloaded_metadata_dir, file_relpath)
return item
# 读取元数据文件
for pa_metadata_table in self._read_metadata(downloaded_metadata_file, metadata_ext=metadata_ext):
for sample in pa_metadata_table.to_pylist():
# 对每个样本应用特征设置函数
for feature_path in feature_paths:
_nested_apply(sample, feature_path, set_feature)
# 生成样本
yield sample_idx, sample
sample_idx += 1
else:
# 如果有过滤器，则将其转换为 PyArrow 表达式
if self.config.filters is not None:
filter_expr = (
pq.filters_to_expression(self.config.filters)
if isinstance(self.config.filters, list)
else self.config.filters
)
for original_file, downloaded_file_or_dir in files:
# 获取下载的文件列表
downloaded_files = [downloaded_file_or_dir] if original_file else downloaded_file_or_dir
for downloaded_file in downloaded_files:
# 获取文件的扩展名
original_file_ext = os.path.splitext(original_file or downloaded_file)[-1]
# 如果文件扩展名不在允许的扩展名列表中，则跳过
if original_file_ext.lower() not in self.EXTENSIONS:
continue
# 构建样本
sample = {self.BASE_COLUMN_NAME: downloaded_file}
# 如果添加标签，则添加标签信息
if add_labels:
sample["label"] = os.path.basename(os.path.dirname(original_file or downloaded_file))
# 如果有过滤器，则对样本进行过滤
if self.config.filters is not None:
pa_table = pa.Table.from_pylist([sample]).filter(filter_expr)
if len(pa_table) == 0:
continue
# 生成样本
yield sample_idx, sample
sample_idx += 1
# 定义一个函数，用于递归地对嵌套数据应用指定的函数
def _nested_apply(item: Any, feature_path: _VisitPath, func: Callable[[Any, _VisitPath], Any]):
# see _visit_with_path() to see how feature paths are constructed
# 对当前项应用指定的函数
item = func(item, feature_path)
if feature_path:
# 获取特征路径的第一个元素
key = feature_path[0]
if key == 0:
# 如果是列表，则递归地对每个元素应用函数
for i in range(len(item)):
item[i] = _nested_apply(item[i], feature_path[1:], func)
else:
# 如果是字典，则递归地对指定键的值应用函数
item[key] = _nested_apply(item[key], feature_path[1:], func)
return item
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
