Big data? 🤗 Datasets vine în ajutor!

În prezent, nu este de neașteptat să te confrunți cu dataseturi de câțiva gigabytes, mai ales dacă planifici să preantreenezi un transformer ca BERT sau GPT-2 de la zero. În aceste cazuri, chiar loading a datelor poate fi o provocare. De exemplu, corpusul WebText folosit pentru a preantreena GPT-2 conține peste 8 milioane de documente și 40 GB de text – încărcarea acestuia în memoria RAM a laptopului tău este probabil să-i facă un atac cardiac!

Norocul este că Datasets 🤗 a fost proiectat pentru a depăși aceste limitări. El te eliberează de problemele de gestionare a memoriei, tratarea dataseturilor ca fișiere memory-mapped, și limitele hard driveului prin streamingul intrărilor dintr-un corpus.

În această secțiune vom explora aceste caracteristici ale 🤗Datasets cu un corpus de 825 GB numit Pile. Să începem!

Ce este Pile?

Pile este un corpus de text englezesc creat de EleutherAI pentru antrenarea large-scale language models. Acesta include o varietate diversă de dataseturi, cuprinzând articole științifice, repositoriuri GitHub și texte web filtrate. Corpusul de antrenare este disponibil în chunkuri de 14 GB, dar în același timp puteți descărca și câteva dintre componenetele individuale. Să începem prin examinarea datasetului PubMed Abstracts, care este un corpus de rezumate din 15 milioane de publicații științifice biomedicale de pe PubMed. Datasetul este în format JSON și este comprimat cu librăria zstandard, așadar prima dată ne trebuie să o instalăm pe aceasta:

!pip install zstandard

În continuare, putem încărca datasetul utilizând metoda pentru fișierele remote pe care am învățat-o în secțiunea 2:

from datasets import load_dataset

# Acest lucru durează câteva minute, așadar poți să te duci să îți iei un ceai sau o cafea între timp :))
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
})

Putem observa că există 15.518.009 de linii și două coloane în datasetul nostru – e foarte mult!

Acum hai să analizăm conținutul primei linii:

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, acesta pare a fi un rezumat dintr-un articol medical. Să vedem cât de mult spațiu RAM am folosit pentru încărcarea datasetului!

Magia memory mappingului

Una dintre modalitățile simple de măsurare a utilizării memoriei în Python este cu biblioteca psutil, care poate fi instalată cu pip:

!pip install psutil

Biblioteca oferă o clasă Process ce ne permite să verificăm utilizarea memoriei procesului curent astfel:

import psutil

# Process.memory_info este exprimat în bytes, așadar convertim la megabyte
print(f"Utilizare RAM: {psutil.Process().memory_info().rss / (1024 * 1024):.2f} MB")

Utilizarea RAM: 5678.33 MB

Aici rss se referă la resident set size , care este partea memoriei ocupată de proces în RAM. Această măsurare include și memoria folosită de Python interpreter și librariile pe care le-am încărcat, așadar cantitatea reală de memorie utilizată pentru încărcarea datelor este puțin mai mică. Pentru comparație, putem să vedem cât de mare este datasetul pe disc folosind atributul dataset_size. Deoarece rezultatul este exprimat în bytes, putem să îl convertim manual la 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

Nice – deși este aproximativ 20 GB, putem încărca și accesa datasetul cu mult mai puțin RAM!

Dacă sunteți familiarizați cu Pandas, rezultatul acesta poate veni ca o surpriză din cauza celebrei rule of thumbă al lui Wes Kinney, care spune că în mod normal aveți nevoie de 5 până la 10 ori mai mult spațiu pe RAM decât mărimea datasetului. Deci 🤗 Datasets această problemă de memory management? 🤗 Datasets tratează fiecare dataset ca un memory-mapped file, care oferă un mapping între spațiul RAM și stocarea pe sistem, ceea ce permite bibliotecii să acceseze și să opereze asupra elementelor datasetului fără a trebui să-l încarce în totalitate în memorie.

Memory-mapped files pot fi și distribuite între mai multe procese, ceea ce permite metodelor cum ar fi Dataset.map() să fie parallelized fără necesitatea mutării sau copierii datasetului. În spatele acestor facilități se află formatul de memorie Apache Arrow și biblioteca pyarrow, care realizează încărcarea datelor și procesarea la viteze fulgerătoare. (Pentru mai multe detalii despre Apache Arrow și compararea sa cu Pandas, vă rugăm să citiți blogul lui Dejan Simic.) Pentru a vedea acest lucru în acțiune, hai să încercăm un speed test prin iterarea asupra tuturor elementelor din datasetul 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'

Aici am folosit modulul timeit al Python pentru a măsura timpul de execuție necesar pentru a rula code_snippet. În mod normal veți putea trece peste un dataset la viteze de câteva sute de MB/s până la câțiva GB/s. Acest lucru funcționează bine pentru majoritatea aplicațiilor, dar uneori veți avea nevoie să lucrați cu un dataset care este prea mare ca să încapă pe hard driveul laptopului tău. De exemplu, dacă am încerca să descarcăm Pile în întregime, am avea nevoie de 825 GB de spațiu liber! Pentru a vă ajuta cu astfel de cazuri, 🤗 Datasets oferă o feature de streaming care permite accesarea și descărcarea elementelor, fără a trebui să descărcați întregul dataset. Hai să vedem cum funcționează!

Streamingul dataseturilor

Pentru a activa dataset streaming este suficient să dați argumentul streaming=True funcției load_dataset(). De exemplu, să încercăm să încărcăm datasetul PubMed Abstracts în streaming mode:

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

În locul Dataset-ului obișnuit pe care l-am întâlnit până acum în acest capitol, obiectul returnat cu streaming=True este un IterableDataset. După nume putem deduce că pentru a accesa elementele dintr-un IterableDataset, trebuie să iterăm prin el. Prin urmare, putem accesa primul element al streamed datased astfel:

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 ...'}

Elementele unui streamed dataset pot fi procesate din mers folosind IterableDataset.map(), ceea ce este util în timpul antrenării dacă aveți nevoie să tokenizeți inputurile. Procesarea se face exact la fel ca și în Capitolul 3, cu singura deosebire fiind că rezultatele sunt returnate una câte una:

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, ...]}

De asemenea, puteți amesteca un streamed dataset utilizând IterableDataset.shuffle(), dar față de Dataset.shuffle() acest lucru va amesteca doar elementele dintr-un buffer_size predefinit:

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 ...'}

În acest exemplu, am selectat un exemplu aleatoriu din primele 10,000 exemple din buffer. Odată ce un exemplu este accesat, locul lui în buffer este completat cu următorul exemplu din corpus (adica exemplarul 10,001 în cazul de mai sus). Puteți selecta elemente dintr-un streamed dataset utilizând funcțiile IterableDataset.take() și IterableDataset.skip(), care acționează în mod similar cu Dataset.select(). De exemplu, pentru a selecta primele 5 exemple din PubMed Abstracts dataset putem face următorul lucru:

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 ...'}]

În mod similar puteți utiliza funcția IterableDataset.skip() pentru a crea splituri de antrenare și validare dintr-un set de date amestecat astfel:

# Săriți primele 1,000 exemple și includeți restul în setul de antrenare
train_dataset = shuffled_dataset.skip(1000)
# Luați primele 1,000 de exemple pentru setul de validare
validation_dataset = shuffled_dataset.take(1000)

Hai să terminăm explorarea streamingului asupra datasetului cu o aplicație comună: combinarea multiplelor dataseturi împreună pentru a crea un singur corpus. 🤗 Datasets oferă o funcție interleave_datasets() care convertește o listă de obiecte IterableDataset într-un singur IterableDataset, unde elementele noului datasetsunt obținute prin alternarea între exemplele sursă. Această funcție este utilă, în special atunci când încercați să combinați dataseturi mari, ca exemplu vom face stream al subsetul FreeLaw din Pile, care reprezintă un dataset de 51 GB de opinii juridice din instanțe din SUA:

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...'}

Această dataset este suficient de mare încât să pună presiune asupra RAM-ului a majorității laptopurilor, dar am reușit să îl încarcăm și să îl accesăm fără să ne facem griji. Acum hai să combinăm exemplele din dataseturile FreeLaw și PubMed Abstracts cu funcția 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...'}]

Aici am folosit funcția islice() din modulul Python itertools pentru a selecta primele două exemple din datasetul combinat, și putem vedea că acestea corespund primelor exemple din fiecare dintre cele două dataseturi originale.

În final, dacă doriți să faceți streaming la Pile în întregime (825 GB), puteți rula toate fișierele pregătite după în acest mod:

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...'}

Acum aveți toate instrumentele necesare pentru a încărca și procesa dataseturi de orice formă și dimensiune – dar, din păcate, va veni un moment în care veți trebui să creați voi înșivă un dataset pentru a rezolva problema pe care o aveți. Acesta este subiectul următoarei secțiuni!