Biblioteca Datasets do Hugging Face: Carregando e Processando Dados

1. Visão Geral

Quando você trabalha com modelos NLP ou faz fine-tune de LLMs, sempre precisa de uma etapa anterior: preparar os dados no formato correto, limpo o suficiente, e conseguir carregá-los no pipeline de treinamento. A biblioteca datasets do Hugging Face resolve exatamente esse problema — ela oferece uma API unificada para carregar datasets do Hub, processar em lote com map/filter, e salvar em um formato otimizado para leitura repetida.

Depois de ler este artigo, você será capaz de carregar datasets do Hugging Face Hub, entender a estrutura Dataset/DatasetDict, saber como fazer pré-processamento de texto com tokenizer, processar datasets grandes com streaming, e escrever testes para validar seu pipeline de processamento de dados.

2. Instalação

pip install datasets transformers

A biblioteca datasets não tem uma dependência obrigatória em transformers, mas a seção de tokenização na seção 6 vai precisar dela.

import datasets
print(datasets.__version__)

3.0.1

3. Carregando datasets do Hugging Face Hub

3.1. load_dataset básico

from datasets import load_dataset

reviews = load_dataset("imdb")
print(reviews)

DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
    test: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
    unsupervised: Dataset({
        features: ['text', 'label'],
        num_rows: 50000
    })
})

load_dataset baixa o script/metadados do dataset imdb do Hub, faz cache do arquivo comprimido localmente (~/.cache/huggingface/datasets), e retorna um DatasetDict — basicamente um dicionário de Datasets divididos por split (train, test, unsupervised).

3.2. Pegando apenas um split específico

train_reviews = load_dataset("imdb", split="train")
print(type(train_reviews), len(train_reviews))

<class 'datasets.arrow_dataset.Dataset'> 25000

Passar split="train" retorna diretamente um Dataset em vez de um DatasetDict, conveniente quando você precisa apenas de um conjunto.

Nota: Se um dataset no Hub requer uma configuração (por exemplo glue), você precisa passar o nome da configuração: load_dataset("glue", "mrpc"). Omitir esse parâmetro vai gerar um erro pedindo para você escolher uma configuração.

4. Explorando a estrutura de um Dataset

sample = train_reviews[0]
print(sample.keys())
print(sample["label"], sample["text"][:60])

dict_keys(['text', 'label'])
0 I rented I AM CURIOUS-YELLOW from my video store because of...

Acessar pelo índice retorna um dicionário Python comum. Dataset também suporta slicing que retorna um dicionário de listas:

batch = train_reviews[0:3]
print(batch["label"])

[0, 0, 0]

Veja o esquema usando features:

print(train_reviews.features)

{'text': Value(dtype='string', id=None),
 'label': ClassLabel(names=['neg', 'pos'], id=None)}

ClassLabel diz a você que 0 corresponde à label neg, 1 é pos — útil quando você quer mapear de volta para nomes de labels.

5. Processando dados com map, filter, select

5.1. filter para filtrar linhas

positive_only = train_reviews.filter(lambda row: row["label"] == 1)
print(len(positive_only))

12500

filter recebe uma função que retorna True/False para cada linha, percorre todo o dataset e mantém apenas as linhas que satisfazem a condição.

5.2. map para transformar dados

def add_length(row):
    row["text_length"] = len(row["text"])
    return row

train_with_length = train_reviews.map(add_length)
print(train_with_length[0]["text_length"])

1640

map aplica uma função a cada linha, retorna um dicionário contendo campos novos (ou atualizando campos existentes) — o resultado é mesclado de volta ao dataset original.

5.3. map em lote para acelerar

def add_length_batched(batch):
    batch["text_length"] = [len(t) for t in batch["text"]]
    return batch

train_fast = train_reviews.map(add_length_batched, batched=True, batch_size=1000)

Com batched=True, a função recebe um dicionário de lote inteiro (cada campo é uma lista) em vez de linha por linha — evita chamar a função Python 25.000 vezes, muito mais rápido com datasets grandes.

5.4. select para pegar um subconjunto

small_train = train_reviews.select(range(100))
print(len(small_train))

100

Use quando você quer testar rapidamente seu pipeline em um subconjunto antes de rodar o dataset completo.

6. Pré-processamento de texto com tokenizer

O caso de uso mais comum na prática: tokenizar texto para passar para o modelo.

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

def tokenize_batch(batch):
    return tokenizer(batch["text"], truncation=True, max_length=256)

tokenized = train_reviews.map(tokenize_batch, batched=True)
print(tokenized.column_names)

['text', 'label', 'input_ids', 'attention_mask']

tokenize_batch recebe um lote e retorna um dicionário com input_ids, attention_mask — map automaticamente mescla esses novos campos no dataset, mantendo os campos antigos. Se você quer remover campos desnecessários (por exemplo, o text bruto para economizar RAM), use remove_columns:

tokenized = train_reviews.map(
    tokenize_batch, batched=True, remove_columns=["text"]
)
print(tokenized.column_names)

['label', 'input_ids', 'attention_mask']

7. Dividindo splits, embaralhando, fragmentando

7.1. Dividindo train/test você mesmo

Se o dataset tem apenas um split, você pode dividir usando train_test_split:

split_result = train_reviews.train_test_split(test_size=0.2, seed=42)
print(split_result)

DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 20000
    })
    test: Dataset({
        features: ['text', 'label'],
        num_rows: 5000
    })
})

seed=42 garante que o resultado da divisão seja reproduzível entre as execuções.

7.2. Embaralhando e fragmentando

shuffled = train_reviews.shuffle(seed=42)
part = shuffled.shard(num_shards=4, index=0)
print(len(part))

6250

shard divide o dataset em num_shards partes iguais — útil quando você quer distribuir dados entre múltiplos workers para treinamento distribuído.

8. Streaming para datasets grandes

Quando o dataset é tão grande que não cabe no disco/RAM, use streaming=True:

streamed = load_dataset("imdb", split="train", streaming=True)
first_three = list(streamed.take(3))
print(len(first_three))

3

Com streaming, datasets não baixa o arquivo inteiro, mas lê registros sequencialmente através de um iterador. Você não pode usar len() nem indexação direta, mas ainda pode fazer map/filter de forma preguiçosa (só aplicadas quando você itera sobre elas).

Nota: Streaming é apropriado para datasets de centenas de GB como conjuntos de pré-treinamento. Com datasets pequenos (algumas centenas de MB), o carregamento normal (streaming=False) é mais rápido porque pode aproveitar o cache e acesso aleatório.

9. Salvando e carregando datasets novamente

Depois de terminar o processamento, salve o resultado para não ter que rodar map novamente:

tokenized.save_to_disk("data/imdb_tokenized")

from datasets import load_from_disk

reloaded = load_from_disk("data/imdb_tokenized")
print(len(reloaded))

25000

save_to_disk salva o dataset em formato Arrow — muito rápido para carregar novamente porque não precisa fazer parse da fonte original. Se você precisa interagir com pandas, converta diretamente:

df = tokenized.select(range(5)).to_pandas()
print(df.columns.tolist())

['label', 'input_ids', 'attention_mask']

10. Escrevendo testes para seu pipeline de processamento de dados

Seu pipeline de processamento de dados deve ter testes para garantir que a lógica de map/filter não quebra silenciosamente quando você modifica o código. Aqui está um teste usando pytest.

import pytest
from datasets import Dataset


def add_length(row):
    row["text_length"] = len(row["text"])
    return row


def test_add_length_computes_correct_value():
    data = Dataset.from_dict({"text": ["hello", "hi there"]})

    result = data.map(add_length)

    assert result[0]["text_length"] == 5
    assert result[1]["text_length"] == 8


def test_filter_keeps_only_positive_labels():
    data = Dataset.from_dict({"label": [0, 1, 1, 0]})

    positive = data.filter(lambda row: row["label"] == 1)

    assert len(positive) == 2
    assert all(row["label"] == 1 for row in positive)

Dataset.from_dict cria um dataset pequeno na memória, sem precisar baixar nada do Hub — perfeito para testes unitários. Os dois testes acima verificam o comportamento correto de map (gera novos campos) e filter (filtra pela condição certa) sem depender de um dataset real na Internet.

pytest test_dataset_pipeline.py -v

test_dataset_pipeline.py::test_add_length_computes_correct_value PASSED
test_dataset_pipeline.py::test_filter_keeps_only_positive_labels PASSED

11. Erros comuns e quando não usar

• Esquecer batched=True ao processar datasets grandes: map sem batch é lento porque chama a função Python para cada linha individual. Com centenas de milhares de linhas, sempre prefira batched=True se a lógica permitir vetorização.
• Usar indexação aleatória em datasets com streaming: streamed[5] vai lançar um erro porque datasets com streaming são iteradores, não suportam acesso aleatório. Use apenas .take(), .skip(), ou iteração sequencial.
• Não definir seed ao usar shuffle/train_test_split: o resultado não será reproduzível entre execuções, dificultando debug ao comparar métricas entre experimentos.
• Cache inchando fora de controle: cada vez que você faz map com uma função diferente, datasets escreve um novo arquivo de cache. Se sua máquina tem pouco espaço em disco, use load_dataset(..., keep_in_memory=True) ou periodicamente limpe ~/.cache/huggingface/datasets.
• Não use datasets para dados relacionais complexos (múltiplas tabelas com joins, precisa de transações). Esta é uma biblioteca otimizada para dados tabulares/texto simples usados em ML, não é um substituto para um banco de dados.

12. Conclusão

A biblioteca datasets oferece uma API limpa para ir de “baixar dataset bruto do Hub” até “dataset tokenizado, pronto para treinamento” em apenas alguns comandos load_dataset, map, filter. Com datasets pequenos, você carrega normalmente e aproveita o cache Arrow; com datasets enormes, mude para streaming=True para não ter que baixar tudo para o disco. Quando estiver construindo um pipeline de processamento de dados para um projeto real, não esqueça de escrever testes para suas funções map/filter usando Dataset.from_dict — isso ajuda você a detectar mais cedo quando a lógica de pré-processamento foi alterada involuntariamente.