Hub documentation

Spark

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Spark

Sparkは、分散環境でのリアルタイムかつ大規模なデータ処理を可能にします。

pyspark_huggingface を使用すると、PySpark で “huggingface” Data Source を介して Hugging Face のデータセットリポジトリにアクセスできます。

Hugging Face Spaces上でSpark Notebooksを試すことで、PySparkとpyspark_huggingfaceがプリインストールされたNotebookを利用できます。

セットアップ

インストール

Hugging Face Datasetsの読み書きを行うには、pyspark_huggingfaceライブラリをインストールする必要があります。

pip install pyspark_huggingface

これにより、認証用のhuggingface_hubや、データセットの読み書き用のpyarrowなど、必要な依存関係もインストールされます。

認証

プライベート/ゲート付きのデータセットリポジトリを読み取ったり、データセットリポジトリに書き込むには、Hugging Face への認証が必要です。

例えば、CLIを以下のように使用できます。

hf auth login

HF_TOKEN 環境変数、またはリーダーに token オプションを渡すことで、Hugging Face トークンを提供することもできます。 認証の詳細については、このガイドを参照してください。

“huggingface” データソースを有効にする

PySpark 4には、カスタムソースからのデータセットを使用できる新しいData Source APIが導入されました。 pyspark_huggingfaceがインストールされている場合、PySparkはこれを自動的にインポートし、“huggingface” Data Sourceを有効にします。

このライブラリは、PySpark 3.5、3.4、3.3 向けの “huggingface” Data Source について、Data Source API をバックポートします。 ただし、この場合、バックポートを有効化し、“huggingface” Data Source を有効にするには、pyspark_huggingface を明示的にインポートする必要があります:

>>> import pyspark_huggingface
huggingface datasource enabled for pyspark 3.x.x (backport from pyspark 4)

読み取る

“huggingface” Data Source は、pyarrow を内部で利用して Arrow データをストリーミングすることで、Hugging Face からデータセットを読み込むことを可能にします。 これは、Parquet データセットなど、Hugging Face 上のサポートされている形式のすべてのデータセットと互換性があります。

例えば、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 サンプルを含む “7M” という名前の設定またはサブセット)を読み込みます。その後、言語ごとの対話数を計算し、データセットをフィルタリングします。

ゲート付きリポジトリにアクセスするためにログインした後、以下を実行できます:

>>> 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 はデータセット内の各データファイルごとに1つのパーティションを持つため、効率的な分散処理が可能になります。

各言語あたりの対話数を計算するために、columns オプションと groupBy() 操作を使用する以下のコードを実行します。 PySpark は Data Source API で predicate push-down を有効にできないため、columns オプションは必要なデータのみを読み込むのに役立ちます。 また、特定の範囲内の値を持つデータのみを読み込むための 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 (string): データセットの分割を選択します(デフォルトは “train”)
  • token (文字列): Hugging Face のトークン

設定や分割を指定する代わりに、読み込むファイルを手動で選択できます:

  • data_dir (文字列): ディレクトリを選択
  • data_files (string): 1 つ以上のファイルを選択します。例: "data/*.parquet" または '["part1.parquet", "par2.parquet"]'

Parquetデータセットの場合:

  • columns (string): 読み込む列のサブセットを選択します。例: '["id"]'
  • filters (string): 条件に一致しないファイルや行グループをスキップするために使用します。例: '[("source", "=", "code_exercises")]'。フィルターは pyarrow.parquet.ParquetDataset に渡されます。

その他のオプションはすべて、https://huggingface.co/docs/datasets/en/package_reference/loading_methods#datasets.load_dataset の [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オプションを指定する必要はありませんが、不要なデータの読み込みを避け、クエリを高速化するために有用です。

書き込み

PySparkのDataframeを「huggingface」Data Sourceを使用して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にデータセットをプッシュする際、2つのモードが利用できます:

  • “overwrite”: データセットが既に存在する場合、上書きします
  • “append”: 既存のデータセットにデータセットを追加します

オプション

write.option()に渡すことができる利用可能なオプションの一覧は以下の通りです:

  • token (文字列): Hugging Face のトークン

ここにさらにオプションを追加する貢献を歓迎します。特にsubsetとsplitです。

ストレージバケット

Storage Bucketsで生データを処理し、DatasetリポジトリにAI対応データを公開する前にそこで実験を行うのが一般的です。

Dataset リポジトリと同様に、buckets/ プレフィックスと data_dir または data_files オプションを指定して Storage Buckets にアクセスします:

>>> 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")
Update on GitHub