Hub documentation
Spark
Spark
Spark 支持在分布式环境中进行实时、大规模的数据处理。
您可以使用 pyspark_huggingface 通过 “huggingface” Data Source 在 PySpark 中访问 Hugging Face 数据集仓库。
在 Hugging Face Spaces 上试用 Spark Notebooks,即可获取预装了 PySpark 和 pyspark_huggingface 的 Notebooks。
设置
安装
要能够读写 Hugging Face Datasets,您需要安装 pyspark_huggingface 库:
pip install pyspark_huggingface这还将安装所需的依赖项,例如用于身份验证的 huggingface_hub 以及用于读写数据集的 pyarrow。
身份验证
您需要对 Hugging Face 进行身份验证,才能读取私有/受限数据集仓库或向您的数据集仓库写入数据。
例如,您可以使用 CLI:
hf auth login也可以通过 HF_TOKEN 环境变量提供您的 Hugging Face 令牌,或者将 token 选项传递给读取器。
有关身份验证的更多详细信息,请参阅此指南。
启用“huggingface”数据源
PySpark 4 引入了新的 Data Source API,允许使用来自自定义数据源的数据集。
如果已安装 pyspark_huggingface,PySpark 会自动导入它并启用 “huggingface” Data Source。
该库还为 PySpark 3.5、3.4 和 3.3 回移了“huggingface”数据源的 Data Source API。
但在这种情况下,需要显式导入 pyspark_huggingface 以激活回移并启用“huggingface”数据源:
>>> import pyspark_huggingface
huggingface datasource enabled for pyspark 3.x.x (backport from pyspark 4)阅读
“huggingface” 数据源允许从 Hugging Face 读取数据集,其底层使用 pyarrow 来流式传输 Arrow 数据。
这与 Hugging Face 上所有 受支持格式 的数据集兼容,例如 Parquet 数据集。
例如,以下是加载 stanfordnlp/imdb 数据集的方法:
>>> import pyspark_huggingface
>>> from pyspark.sql import SparkSession
>>> spark = SparkSession.builder.appName("demo").getOrCreate()
>>> df = spark.read.format("huggingface").load("stanfordnlp/imdb")这里还有一个使用 BAAI/Infinity-Instruct 数据集的示例。 这是一个受控仓库,用户必须先接受使用条款才能访问。 它还有多个子集,即“3M”和“7M”。因此我们需要指定要加载哪一个。

我们使用 .format() 函数来使用 “huggingface” 数据源,并使用 .load() 加载数据集(更准确地说,是加载名为 “7M” 的配置或子集,其中包含 700 万个样本)。然后我们计算每种语言的对话数量,并过滤数据集。
登录以访问受限制仓库后,我们可以运行:
>>> import pyspark_huggingface
>>> from pyspark.sql import SparkSession
>>> spark = SparkSession.builder.appName("demo").getOrCreate()
>>> df = spark.read.format("huggingface").option("config", "7M").load("BAAI/Infinity-Instruct")
>>> df.show()
+---+----------------------------+-----+----------+--------------------+
| id| conversations|label|langdetect| source|
+---+----------------------------+-----+----------+--------------------+
| 0| [{human, def exti...| | en| code_exercises|
| 1| [{human, See the ...| | en| flan|
| 2| [{human, This is ...| | en| flan|
| 3| [{human, If you d...| | en| flan|
| 4| [{human, In a Uni...| | en| flan|
| 5| [{human, Read the...| | en| flan|
| 6| [{human, You are ...| | en| code_bagel|
| 7| [{human, I want y...| | en| Subjective|
| 8| [{human, Given th...| | en| flan|
| 9|[{human, 因果联系原则是法...| | zh-cn| Subjective|
| 10| [{human, Provide ...| | en|self-oss-instruct...|
| 11| [{human, The univ...| | en| flan|
| 12| [{human, Q: I am ...| | en| flan|
| 13| [{human, What is ...| | en| OpenHermes-2.5|
| 14| [{human, In react...| | en| flan|
| 15| [{human, Write Py...| | en| code_exercises|
| 16| [{human, Find the...| | en| MetaMath|
| 17| [{human, Three of...| | en| MetaMath|
| 18| [{human, Chandra ...| | en| MetaMath|
| 19|[{human, 用经济学知识分析...| | zh-cn| Subjective|
+---+----------------------------+-----+----------+--------------------+这将以流式方式加载数据集,输出的 DataFrame 中每个数据文件对应一个分区,以实现高效的分布式处理。
要计算每种语言的对话数量,我们运行了这段代码,其中使用了 columns 选项和 groupBy() 操作。
columns 选项非常有用,因为它可以只加载我们所需的数据,因为 PySpark 在使用 Data Source API 时不支持谓词下推。
此外,还有一个 filters 选项,用于只加载值在特定范围内的数据。
>>> df_langdetect_only = (
... spark.read.format("huggingface")
... .option("config", "7M")
... .option("columns", '["langdetect"]')
... .load("BAAI/Infinity-Instruct")
... )
>>> df_langdetect_only.groupBy("langdetect").count().show()
+----------+-------+
|langdetect| count|
+----------+-------+
| en|6697793|
| zh-cn| 751313|
+----------+-------+要过滤数据集并仅保留中文对话:
>>> df_chinese_only = (
... spark.read.format("huggingface")
... .option("config", "7M")
... .option("filters", '[("langdetect", "=", "zh-cn")]')
... .load("BAAI/Infinity-Instruct")
... )
>>> df_chinese_only.show()
+---+----------------------------+-----+----------+----------+
| id| conversations|label|langdetect| source|
+---+----------------------------+-----+----------+----------+
| 9|[{human, 因果联系原则是法...| | zh-cn|Subjective|
| 19|[{human, 用经济学知识分析...| | zh-cn|Subjective|
| 38| [{human, 某个考试共有A、...| | zh-cn|Subjective|
| 39|[{human, 撰写一篇关于斐波...| | zh-cn|Subjective|
| 57|[{human, 总结世界历史上的...| | zh-cn|Subjective|
| 61|[{human, 生成一则广告词。...| | zh-cn|Subjective|
| 66|[{human, 描述一个有效的团...| | zh-cn|Subjective|
| 94|[{human, 如果比利和蒂芙尼...| | zh-cn|Subjective|
|102|[{human, 生成一句英文名言...| | zh-cn|Subjective|
|106|[{human, 写一封感谢信,感...| | zh-cn|Subjective|
|118| [{human, 生成一个故事。}...| | zh-cn|Subjective|
|174|[{human, 高胆固醇水平的后...| | zh-cn|Subjective|
|180|[{human, 基于以下角色信息...| | zh-cn|Subjective|
|192|[{human, 请写一篇文章,概...| | zh-cn|Subjective|
|221|[{human, 以诗歌形式表达对...| | zh-cn|Subjective|
|228|[{human, 根据给定的指令,...| | zh-cn|Subjective|
|236|[{human, 打开一个新的生成...| | zh-cn|Subjective|
|260|[{human, 生成一个有关未来...| | zh-cn|Subjective|
|268|[{human, 如果有一定数量的...| | zh-cn|Subjective|
|273| [{human, 题目:小明有5个...| | zh-cn|Subjective|
+---+----------------------------+-----+----------+----------+也可以在加载后的 DataFrame 上应用过滤器或移除列,但在加载时执行这些操作效率更高,尤其是在处理 Parquet 数据集时。 实际上,Parquet 在文件和行组级别包含元数据,这使得可以跳过不包含满足条件样本的数据集部分。Parquet 中的列也可以独立加载,从而允许跳过被排除的列,避免加载不必要的数据。
选项
以下是可以传递给 read..option() 的可用选项列表:
config(字符串):选择数据集的子集/配置split(字符串):选择数据集的拆分(默认为“train”)token(字符串):您的 Hugging Face 令牌
而不是指定配置或拆分,您可以手动选择要加载的文件:
data_dir(字符串):选择一个目录data_files(字符串):选择一个或多个文件,例如"data/*.parquet"或'["part1.parquet", "par2.parquet"]'
对于 Parquet 数据集:
columns(字符串):选择要加载的列子集,例如'["id"]'filters(字符串):用于跳过不符合条件的文件和行组,例如'[("source", "=", "code_exercises")]'。过滤器会传递给 pyarrow.parquet.ParquetDataset。
其他任何选项都会作为参数传递给 [datasets.load_dataset] (https://huggingface.co/docs/datasets/en/package_reference/loading_methods#datasets.load_dataset)
运行 SQL 查询
准备好 PySpark Dataframe 后,您可以使用 spark.sql 运行 SQL 查询:
>>> import pyspark_huggingface
>>> from pyspark.sql import SparkSession
>>> spark = SparkSession.builder.appName("demo").getOrCreate()
>>> df = (
... spark.read.format("huggingface")
... .option("config", "7M")
... .option("columns", '["source"]')
... .load("BAAI/Infinity-Instruct")
... )
>>> spark.sql("SELECT source, count(*) AS total FROM {df} GROUP BY source ORDER BY total DESC", df=df).show()
+--------------------+-------+
| source| total|
+--------------------+-------+
| flan|2435840|
| Subjective|1342427|
| OpenHermes-2.5| 855478|
| MetaMath| 690138|
| code_exercises| 590958|
|Orca-math-word-pr...| 398168|
| code_bagel| 386649|
| MathInstruct| 329254|
|python-code-datas...| 88632|
|instructional_cod...| 82920|
| CodeFeedback| 79513|
|self-oss-instruct...| 50467|
|Evol-Instruct-Cod...| 43354|
|CodeExercise-Pyth...| 27159|
|code_instructions...| 23130|
| Code-Instruct-700k| 10860|
|Glaive-code-assis...| 9281|
|python_code_instr...| 2581|
|Python-Code-23k-S...| 2297|
+--------------------+-------+同样,指定 columns 选项并非必需,但有助于避免加载不必要的数据并加快查询速度。
写入
您可以使用“huggingface”数据源将 PySpark Dataframe 写入 Hugging Face。 它会以分布式方式并行上传 Parquet 文件,并且只有在所有文件都上传完成后才会提交这些文件。 其工作原理如下:
>>> import pyspark_huggingface
>>> df.write.format("huggingface").save("username/dataset_name")以下是如何使用该函数将 BAAI/Infinity-Instruct 数据集的过滤版本写回 Hugging Face 的方法。
首先,您需要创建一个数据集仓库,例如 username/Infinity-Instruct-Chinese-Only(如果需要,可以将其设置为私有)。
然后,确保您已完成身份验证并且可以使用“huggingface”数据源,将mode设置为“overwrite”(如果要扩展现有数据集,则设置为“append”),并使用.save()推送到 Hugging Face:
>>> df_chinese_only.write.format("huggingface").mode("overwrite").save("username/Infinity-Instruct-Chinese-Only")

模式
将数据集推送到 Hugging Face 时,有两种可用模式:
- “overwrite”:如果数据集已存在,则覆盖
- “append”:将数据集追加到现有数据集
选项
以下是可以传递给 write.option() 的可用选项列表:
token(字符串):您的 Hugging Face 令牌
欢迎贡献更多选项,特别是 subset 和 split。
存储桶
通常,人们会在Storage Buckets中处理原始数据并进行实验,然后再将可用于 AI 的数据发布到 Dataset 仓库中。
访问存储桶的方式与访问数据集仓库相同,但需使用 buckets/ 前缀,并配合 data_dir 或 data_files 选项:
>>> df = spark.read.format("huggingface").option("data_dir", "data").load("buckets/username/my-bucket")
>>> # OR with a glob pattern
>>> # df = spark.read.format("huggingface").option("data_files", "data/*.parquet").load("buckets/username/my-bucket")
>>> df.write.format("huggingface").option("data_dir", "new-data").save("buckets/username/my-bucket")