Big data? Ci pensa 🤗 Datasets!

Al giorno d’oggi non è raro trovarsi a lavorare con dataset grandi diversi gigabyte, soprattutto quando si vuole addestrare un transformer come BERT o GPT-2 da zero. In questi casi, persino caricare i dati può essere un’impresa difficile. Ad esempio, il corpus WebText utilizzato per preaddestrare GPT-2 contiente più di 8 milioni di documenti e 40gb di testo — caricare un dataset del genere sulla RAM del tuo portatile gli farebbe venire un colpo!

Per fortuna, 🤗 Datasets è stato sviluppato per superare queste limitazioni, e può risolvere i problemi relativi alla gestione della memoria trattando i dataset come file memory-mapped, e quelli relativi ai limiti del disco rigido attraverso lo stream processing delle voci del corpus.

In questa sezione esploreremo queste funzionalità di 🤗 Datasets con un enorme corpus di 825 GB conosciuto come Pile. Iniziamo!

Cos’è Pile?

The Pile è un corpus testuale creato da EleutherAI per addestrare modelli di linguaggio su grande scala. Include un grande varietà di dataset, a partire da articoli scientifici, repository di codici da GitHub, e testi dal web filtrati. Il corpus di addestramento è disponibili in frammenti da 14 GB, ed è possibile scaricare diverse delle componenti singole. Iniziamo dando uno sguardo al dataset PubMed Abstracts, un corpus di abstract da 15 milioni di pubblicazioni in ambito biomedico da PubMed. Il dataset è in formato JSON Lines ed è stato compressato usando la libreria zstandard, per cui dobbiamo prima installarla:

!pip install zstandard

Ora, possiamo caricare il dataset utilizzando il meotodo per file remoti che abbiamo visto nella sezione 2:

from datasets import load_dataset

# Ci vuole qualche minuto per l'esecuzione, quindi preparati un tè o un caffè nell'attesa :)
data_files = "https://the-eye.eu/public/AI/pile_preliminary_components/PUBMED_title_abstracts_2019_baseline.jsonl.zst"
pubmed_dataset = load_dataset("json", data_files=data_files, split="train")
pubmed_dataset

Dataset({
    features: ['meta', 'text'],
    num_rows: 15518009
})

Possiamo vedere che ci sono 15.518.009 righe e 2 colonne nel nostro dataset — un bel po’!

Ispezioniamo i contenuti del primo esempio:

pubmed_dataset[0]

{'meta': {'pmid': 11409574, 'language': 'eng'},
 'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection.\nTo determine the prevalence of hypoxaemia in children aged under 5 years suffering acute lower respiratory infections (ALRI), the risk factors for hypoxaemia in children under 5 years of age with ALRI, and the association of hypoxaemia with an increased risk of dying in children of the same age ...'}

Okay, questo sembra proprio l’abstract di un articolo di medicina. Ora vediamo quanta RAM è stata usata per caricare il dataset!

La magia del memory mapping

Un modo semplice per calcolare l’uso di memoria su Python è utilizzando la libreria psutil, che può essere installata con pip come segue:

!pip install psutil

psutil offre una classe Process che permette di controllare l’utilizzo della memoria del processo attuale come segue::

import psutil

# Process.memory_info mostra i dati in byte, quindi convertiamo in megabyte
print(f"RAM used: {psutil.Process().memory_info().rss / (1024 * 1024):.2f} MB")

RAM used: 5678.33 MB

L’attributo rss qui fa riferimento alla grandezza del resident set, che equivale alla frazione di memoria che il processo occupa nella RAM. Questo valore include inoltre la memoria utilizzata dall’interprete Python e dalle librerie caricate, per cui l’ammontare effettivo utilizzato per caricare il dataset è un po’ più piccolo. Per fare un confronto, vediamo quant’è grande il dataset su disco utilizzando l’attributo dataset_size. Come prima, il risultato è espresso in byte, e abbiamo bisogno di convertirlo in gigabyte:

print(f"Number of files in dataset : {pubmed_dataset.dataset_size}")
size_gb = pubmed_dataset.dataset_size / (1024**3)
print(f"Dataset size (cache file) : {size_gb:.2f} GB")

Number of files in dataset : 20979437051
Dataset size (cache file) : 19.54 GB

Bene — nonostante sia grande quasi 30 GB, siamo in grado di caricare e accedere al dataset utilizzando molta meno RAM!

Se hai dimestichezza con Pandas, questo risultato potrebbe sorprenderti, vista la famosa regola di Wes Kinney, ovvero che, in linea di massima, serve una RAM 5-10 volte più grande del dataset che vuoi caricare. Come fa 🤗 Datasets a risolvere questo problema di gestione della memoria? 🤗 Datasets tratta ogni dataset come un file mappato in memoria, il che permette di avere un mapping tra la RAM e l’archiviazione dei file di sistema, che permette alla librera di accedere e operare su elementi del dataset senza doverli caricare completamente in memoria.

I file mappati in memoria possono inoltre essre condivisi su più processi, il che permette a metodi come Dataset.map() di poter essere eseguiti in parallelo senza bisogno di spostare o copiare il dataset. Dietro le quinte, tutto ciò è realizzato dal formato di memoria Apache Arrow e dalla libreria pyarrow, che rendono più veloci il caricamento e il processamento dei dati. (per maggiori dettagli su Apache Arrow, e per un confronto con Pandas, dai un’occhiata al post di Dejan Simic.) Per vederlo in azione, eseguiamo un piccolo test di velocità con un loop su tutti gli elementi nel dataset PubMed Abstracts:

import timeit

code_snippet = """batch_size = 1000

for idx in range(0, len(pubmed_dataset), batch_size):
    _ = pubmed_dataset[idx:idx + batch_size]
"""

time = timeit.timeit(stmt=code_snippet, number=1, globals=globals())
print(
    f"Iterated over {len(pubmed_dataset)} examples (about {size_gb:.1f} GB) in "
    f"{time:.1f}s, i.e. {size_gb/time:.3f} GB/s"
)

'Iterated over 15518009 examples (about 19.5 GB) in 64.2s, i.e. 0.304 GB/s'

Abbiamo usato il modulo di Python timeit per calcolare il tempo di esecuzione impiegato da code_snippet. Tipicamente l’iterazione su un dataset impiega un tempo che va da un decimo di GB al secondo, a diversi GB al secondo. Questo funziona perfettamente per la maggior parte delle applicazioni, ma a volte avrai bisogno di lavorare con un dataset che è troppo grande persino per essere salvato sul tuo portatile. Ad esempio, se cercassimo di scaricare Pile per intero, avremo bisogno di 825 GB di spazio libero su disko! In questi casi, 🤗 Datasets permette di utilizzare processi di streaming che ci permettono di scaricare e accedere al volo ai dati, senza bisogno di scaricare l’intero dataset. Diamo un’occhiata a come funziona.

Streaming di dataset

Per abilitare lo streaming dei dataset devi semplicemente passare l’argomento streaming=True alla funzione load_dataset(). Ad esempio, carichiamo un’altra volta il dataset PubMed Abstract, ma in modalità streaming:

pubmed_dataset_streamed = load_dataset(
    "json", data_files=data_files, split="train", streaming=True
)

Invece del solito Dataset che abbiamo incontrato in precedenza in questo capitolo, l’oggetto ritornato con streaming=True' è un IterableDataset. Come suggerito dal nome, per accedere agli elementi di un IterableDataset`, dobbiamo iterare di esso. Possiamo accedere al primo elemento del nostro dataset in streaming come segue:

next(iter(pubmed_dataset_streamed))

{'meta': {'pmid': 11409574, 'language': 'eng'},
 'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection.\nTo determine the prevalence of hypoxaemia in children aged under 5 years suffering acute lower respiratory infections (ALRI), the risk factors for hypoxaemia in children under 5 years of age with ALRI, and the association of hypoxaemia with an increased risk of dying in children of the same age ...'}

Gli elementi di un dataset in streaming possono essere processati al volo utilizzando IterableDataset.map(), che è utile durante l’addestramento se hai bisogno di tokenizzare gli input. Il processo è uguale a quello che abbiamo utilizzato per tokenizzare il nostro dataset nel Capitolo 3, con l’unica differenza che ora ritorneremo gli output uno alla volta:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
tokenized_dataset = pubmed_dataset_streamed.map(lambda x: tokenizer(x["text"]))
next(iter(tokenized_dataset))

{'input_ids': [101, 4958, 5178, 4328, 6779, ...], 'attention_mask': [1, 1, 1, 1, 1, ...]}

È anche possibile mescolare un dataset in streaming utilizzato Iterabledataset.shuffle(), ma a differenza di Dataset.shuffle(), questo metodo mescola solo gli elementi in un buffer_size predefinito:

shuffled_dataset = pubmed_dataset_streamed.shuffle(buffer_size=10_000, seed=42)
next(iter(shuffled_dataset))

{'meta': {'pmid': 11410799, 'language': 'eng'},
 'text': 'Randomized study of dose or schedule modification of granulocyte colony-stimulating factor in platinum-based chemotherapy for elderly patients with lung cancer ...'}

In questo esempio, abbiamo selezionato un esempio casuale dai primi 10.000 esempi nel buffer. Una volta che accediamo a un esempio, il suo posto nel buffer è subito occupato dall’esempio successivo nel corpus (in questo caso l’esempio 10.0001). Puoi inoltre selezionare gli elementi da un dataset in streaming utilizzando le funzioni IterableDataset.take() a IterableDataset.skip(), che funzionano un po’ come Dataset.select(). Ad esempio, per selezionare i primi 5 esempi nel dataset PubMed Abstract dovremmo fare come segue:

dataset_head = pubmed_dataset_streamed.take(5)
list(dataset_head)

[{'meta': {'pmid': 11409574, 'language': 'eng'},
  'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection ...'},
 {'meta': {'pmid': 11409575, 'language': 'eng'},
  'text': 'Clinical signs of hypoxaemia in children with acute lower respiratory infection: indicators of oxygen therapy ...'},
 {'meta': {'pmid': 11409576, 'language': 'eng'},
  'text': "Hypoxaemia in children with severe pneumonia in Papua New Guinea ..."},
 {'meta': {'pmid': 11409577, 'language': 'eng'},
  'text': 'Oxygen concentrators and cylinders ...'},
 {'meta': {'pmid': 11409578, 'language': 'eng'},
  'text': 'Oxygen supply in rural africa: a personal experience ...'}]

Allo stesso modo, è possibile utilizzare la funzione IterableDataset.skip() per creare sezioni di addestramento e di validazione da un dataset mescolato, come segue:

# Salta i primi 1.000 esempi, il resto viene incluso nell'insieme di addestramento
train_dataset = shuffled_dataset.skip(1000)
# Includi i primi 1.000 esempi nell'insieme di validazione
validation_dataset = shuffled_dataset.take(1000)

Concludiamo la nostra ricognizione dello streaming di dataset con un’applicazione comune: la combinazione di più dataset per creare un unico corpus. 🤗 Datasets fornisce una funzione interleave_datasets(), che converte una lista di oggetti IterableDataset in un unico IterableDataset, dove gli elementi del nuovo dataset sono ottenuti alternando tra gli esempi forniti. Questa funzione è particolarmente utile quando cerchiamo di combinare dataset di grandi dimensioni, come esempio possiamo utilizzare in streaming la sezione FreeLaw del Pile, un dataset di 51 GB di pareri legali dai tribunali statunitensi:

law_dataset_streamed = load_dataset(
    "json",
    data_files="https://the-eye.eu/public/AI/pile_preliminary_components/FreeLaw_Opinions.jsonl.zst",
    split="train",
    streaming=True,
)
next(iter(law_dataset_streamed))

{'meta': {'case_ID': '110921.json',
  'case_jurisdiction': 'scotus.tar.gz',
  'date_created': '2010-04-28T17:12:49Z'},
 'text': '\n461 U.S. 238 (1983)\nOLIM ET AL.\nv.\nWAKINEKONA\nNo. 81-1581.\nSupreme Court of United States.\nArgued January 19, 1983.\nDecided April 26, 1983.\nCERTIORARI TO THE UNITED STATES COURT OF APPEALS FOR THE NINTH CIRCUIT\n*239 Michael A. Lilly, First Deputy Attorney General of Hawaii, argued the cause for petitioners. With him on the brief was James H. Dannenberg, Deputy Attorney General...'}

Questo dataset è abbastanza grande da mettere sotto sforzo la RAM di molto portatili, ma siamo riusciti a caricarlo e accedervi senza alcun problema! Ora cominiamo gli esempi di FreeLaw e di PubMed Abstracts con la funzione interleave_datasets():

from itertools import islice
from datasets import interleave_datasets

combined_dataset = interleave_datasets([pubmed_dataset_streamed, law_dataset_streamed])
list(islice(combined_dataset, 2))

[{'meta': {'pmid': 11409574, 'language': 'eng'},
  'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection ...'},
 {'meta': {'case_ID': '110921.json',
   'case_jurisdiction': 'scotus.tar.gz',
   'date_created': '2010-04-28T17:12:49Z'},
  'text': '\n461 U.S. 238 (1983)\nOLIM ET AL.\nv.\nWAKINEKONA\nNo. 81-1581.\nSupreme Court of United States.\nArgued January 19, 1983.\nDecided April 26, 1983.\nCERTIORARI TO THE UNITED STATES COURT OF APPEALS FOR THE NINTH CIRCUIT\n*239 Michael A. Lilly, First Deputy Attorney General of Hawaii, argued the cause for petitioners. With him on the brief was James H. Dannenberg, Deputy Attorney General...'}]

Abbiamo utilizzato la funzione islice() del modulo Python itertools per selezionare i primi due esempi dai dataset combinati, e abbiamo visto che corrispondono ai primi esempi di ognuno dei due dataset originali.

Infine, se vuoi processare il Pile in streaming, in tutti i suoi 825 GB, puoi recuperare tutti i file preparati, come segue:

base_url = "https://the-eye.eu/public/AI/pile/"
data_files = {
    "train": [base_url + "train/" + f"{idx:02d}.jsonl.zst" for idx in range(30)],
    "validation": base_url + "val.jsonl.zst",
    "test": base_url + "test.jsonl.zst",
}
pile_dataset = load_dataset("json", data_files=data_files, streaming=True)
next(iter(pile_dataset["train"]))

{'meta': {'pile_set_name': 'Pile-CC'},
 'text': 'It is done, and submitted. You can play “Survival of the Tastiest” on Android, and on the web...'}

Ora hai a tua disposizione tutti gli strumenti per caricare e processare dataset di ogni tipo — ma a meno che tu non sia estremamente fortunato, arriverà un momento nel tuo cammino in cui dovrai effettivamente creare un dataset per risolvere i tuoi problemi. Questo sarà argomento della prossima sezione!