Spark

O Spark permite o processamento de dados em tempo real e em larga escala em um ambiente distribuído.

Você pode usar pyspark_huggingface para acessar repositórios de datasets do Hugging Face no PySpark por meio da Data Source “huggingface”.

Experimente os Spark Notebooks nos Hugging Face Spaces para obter Notebooks com PySpark e pyspark_huggingface pré-instalados.

Configuração

Instalação

Para poder ler e gravar no Hugging Face Datasets, você precisa instalar a biblioteca pyspark_huggingface:

pip install pyspark_huggingface

Isso também instalará as dependências necessárias, como huggingface_hub para autenticação e pyarrow para leitura e gravação de conjuntos de dados.

Autenticação

Você precisa autenticar-se no Hugging Face para ler repositórios de datasets privados/limitados ou para gravar em seus repositórios de datasets.

Você pode usar a CLI, por exemplo:

hf auth login

Também é possível fornecer seu token do Hugging Face com a variável de ambiente HF_TOKEN ou passando a opção token para o leitor. Para mais detalhes sobre autenticação, confira este guia.

Ative a fonte de dados “huggingface”

O PySpark 4 veio com uma nova API de Data Source que permite usar conjuntos de dados de fontes personalizadas. Se pyspark_huggingface estiver instalado, o PySpark o importará automaticamente e habilitará a Data Source “huggingface”.

A biblioteca também faz o backport da API de Data Source para o Data Source “huggingface” para o PySpark 3.5, 3.4 e 3.3. No entanto, neste caso, pyspark_huggingface deve ser importado explicitamente para ativar o backport e habilitar o Data Source “huggingface”:

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

Ler

A Fonte de Dados “huggingface” permite ler conjuntos de dados do Hugging Face, usando pyarrow por baixo para transmitir dados Arrow. Isso é compatível com todos os conjuntos de dados em formato suportado no Hugging Face, como conjuntos de dados Parquet.

Por exemplo, aqui está como carregar o conjunto de dados stanfordnlp/imdb:

>>> import pyspark_huggingface
>>> from pyspark.sql import SparkSession
>>> spark = SparkSession.builder.appName("demo").getOrCreate()
>>> df = spark.read.format("huggingface").load("stanfordnlp/imdb")

Aqui está outro exemplo com o conjunto de dados BAAI/Infinity-Instruct. Trata-se de um repositório restrito, os usuários precisam aceitar os termos de uso antes de acessá-lo. Ele também possui vários subconjuntos, a saber, “3M” e “7M”. Portanto, precisamos especificar qual deles carregar.

Usamos a função .format() para usar a Data Source “huggingface” e .load() para carregar o dataset (mais precisamente, a config ou subconjunto chamado “7M”, contendo 7M amostras). Em seguida, calculamos o número de diálogos por idioma e filtramos o dataset.

Após fazer login para acessar o repositório restrito, podemos executar:

>>> 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|
+---+----------------------------+-----+----------+--------------------+

Isso carrega o dataset de forma em streaming, e o DataFrame de saída possui uma partição por arquivo de dados no dataset para permitir um processamento distribuído eficiente.

Para calcular o número de diálogos por idioma, executamos este código que usa a opção columns e uma operação groupBy(). A opção columns é útil para carregar apenas os dados de que precisamos, já que o PySpark não habilita o predicate push-down com a Data Source API. Também existe a opção filters para carregar apenas dados com valores dentro de uma determinada faixa.

>>> 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|
+----------+-------+

Para filtrar o conjunto de dados e manter apenas os diálogos em chinês:

>>> 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|
+---+----------------------------+-----+----------+----------+

Também é possível aplicar filtros ou remover colunas no DataFrame carregado, mas é mais eficiente fazer isso durante o carregamento, especialmente em datasets Parquet. De fato, o Parquet contém metadados no nível do arquivo e do grupo de linhas, o que permite pular partes inteiras do dataset que não contêm amostras que atendem aos critérios. As colunas no Parquet também podem ser carregadas de forma independente, o que permite pular as colunas excluídas e evitar o carregamento de dados desnecessários.

Opções

Aqui está a lista de opções disponíveis que você pode passar para read..option():

• config (string): seleciona um subconjunto/configuração do conjunto de dados
• split (string): seleciona um split do dataset (o padrão é “train”)
• token (string): seu token do Hugging Face

Em vez de especificar um config ou split, você pode selecionar manualmente quais arquivos carregar:

• data_dir (string): selecione um diretório
• data_files (string): selecione um ou vários arquivos, por exemplo, "data/*.parquet" ou '["part1.parquet", "par2.parquet"]'

Para conjuntos de dados Parquet:

• columns (string): seleciona um subconjunto de colunas para carregar, por exemplo, '["id"]'
• filters (string): para ignorar arquivos e grupos de linhas que não correspondem a um critério, por exemplo, '[("source", "=", "code_exercises")]'. Os filtros são passados para pyarrow.parquet.ParquetDataset.

Qualquer outra opção é passada como argumento para [datasets.load_dataset] (https://huggingface.co/docs/datasets/en/package_reference/loading_methods#datasets.load_dataset)

Executar consultas SQL

Depois de ter seu Dataframe do PySpark pronto, você pode executar consultas SQL usando spark.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|
+--------------------+-------+

Novamente, especificar a opção columns não é necessário, mas é útil para evitar o carregamento de dados desnecessários e tornar a consulta mais rápida.

Escrever

Você pode gravar um Dataframe do PySpark no Hugging Face usando a Data Source “huggingface”. Ele faz o upload de arquivos Parquet em paralelo de forma distribuída e só confirma os arquivos depois que todos forem enviados. Funciona da seguinte forma:

>>> import pyspark_huggingface
>>> df.write.format("huggingface").save("username/dataset_name")

Aqui está como podemos usar essa função para gravar a versão filtrada do conjunto de dados BAAI/Infinity-Instruct de volta no Hugging Face.

Primeiro, você precisa criar um repositório de dataset, por exemplo, username/Infinity-Instruct-Chinese-Only (você pode configurá-lo como privado, se desejar). Em seguida, verifique se você está autenticado e pode usar a fonte de dados “huggingface”, defina o mode como “overwrite” (ou “append”, se quiser estender um dataset existente) e envie para o Hugging Face com .save():

>>> df_chinese_only.write.format("huggingface").mode("overwrite").save("username/Infinity-Instruct-Chinese-Only")

Modo

Dois modos estão disponíveis ao enviar um dataset para o Hugging Face:

• “overwrite”: substitui o conjunto de dados se ele já existir
• “append”: adiciona o conjunto de dados a um conjunto de dados existente

Opções

Aqui está a lista de opções disponíveis que você pode passar para write.option():

• token (string): seu token do Hugging Face

Contribuições são bem-vindas para adicionar mais opções aqui, em particular subset e split.

Buckets de Armazenamento

É comum processar dados brutos em Storage Buckets e realizar experimentos ali antes de publicar dados prontos para IA em repositórios de Dataset.

Acesse os Storage Buckets da mesma forma que os repositórios de Dataset, mas com o prefixo buckets/ e com as opções data_dir ou 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")