Machine learning dataset loaders for testing and examples

Loaders for various machine learning datasets for testing and example scripts. Previously in thinc.extra.datasets.

Setup and installation

The package can be installed via pip:

pip install ml-datasets

Loaders

Loaders can be imported directly or used via their string name (which is useful if they're set via command line arguments). Some loaders may take arguments – see the source for details.

# Import directlyfromml_datasetsimportimdbtrain_data, dev_data=imdb()

# Load via registryfromml_datasetsimportloadersimdb_loader=loaders.get("imdb")
train_data, dev_data=imdb_loader()

Available loaders

NLP datasets

ID / Function	Description	NLP task	From URL
imdb	IMDB sentiment dataset	Binary classification: sentiment analysis	✓
dbpedia	DBPedia ontology dataset	Multi-class single-label classification	✓
cmu	CMU movie genres dataset	Multi-class, multi-label classification	✓
quora_questions	Duplicate Quora questions dataset	Detecting duplicate questions	✓
reuters	Reuters dataset (texts not included)	Multi-class multi-label classification	✓
snli	Stanford Natural Language Inference corpus	Recognizing textual entailment	✓
stack_exchange	Stack Exchange dataset	Question Answering
ud_ancora_pos_tags	Universal Dependencies Spanish AnCora corpus	POS tagging	✓
ud_ewtb_pos_tags	Universal Dependencies English EWT corpus	POS tagging	✓
wikiner	WikiNER data	Named entity recognition

Other ML datasets

ID / Function	Description	ML task	From URL
mnist	MNIST data	Image recognition	✓

Dataset details

IMDB

Each instance contains the text of a movie review, and a sentiment expressed as 0 or 1.

train_data, dev_data=ml_datasets.imdb()
fortext, annotintrain_data[0:5]:
print(f"Review: {text}")
print(f"Sentiment: {annot}")

• Download URL: http://ai.stanford.edu/~amaas/data/sentiment/
• Citation: Andrew L. Maas et al., 2011

Property	Training	Dev
# Instances	25000	25000
Label values	{0, 1}	{0, 1}
Labels per instance	Single	Single
Label distribution	Balanced (50/50)	Balanced (50/50)

DBPedia

Each instance contains an ontological description, and a classification into one of the 14 distinct labels.

train_data, dev_data=ml_datasets.dbpedia()
fortext, annotintrain_data[0:5]:
print(f"Text: {text}")
print(f"Category: {annot}")

• Download URL: Via fast.ai
• Original citation: Xiang Zhang et al., 2015

Property	Training	Dev
# Instances	560000	70000
Label values	1-14	1-14
Labels per instance	Single	Single
Label distribution	Balanced	Balanced

CMU

Each instance contains a movie description, and a classification into a list of appropriate genres.

train_data, dev_data=ml_datasets.cmu()
fortext, annotintrain_data[0:5]:
print(f"Text: {text}")
print(f"Genres: {annot}")

• Download URL: http://www.cs.cmu.edu/~ark/personas/
• Original citation: David Bamman et al., 2013

Property	Training	Dev
# Instances	41793	0
Label values	363 different genres	-
Labels per instance	Multiple	-
Label distribution	Imbalanced: 147 labels with less than 20 examples, while Drama occurs more than 19000 times	-

Quora

train_data, dev_data=ml_datasets.quora_questions()
forquestions, annotintrain_data[0:50]:
q1, q2=questionsprint(f"Question 1: {q1}")
print(f"Question 2: {q2}")
print(f"Similarity: {annot}")

Each instance contains two quora questions, and a label indicating whether or not they are duplicates (0: no, 1: yes). The ground-truth labels contain some amount of noise: they are not guaranteed to be perfect.

• Download URL: http://qim.fs.quoracdn.net/quora_duplicate_questions.tsv
• Original citation: Kornél Csernai et al., 2017

Property	Training	Dev
# Instances	363859	40429
Label values	{0, 1}	{0, 1}
Labels per instance	Single	Single
Label distribution	Imbalanced: 63% label 0	Imbalanced: 63% label 0

Registering loaders

Loaders can be registered externally using the loaders registry as a decorator. For example:

@ml_datasets.loaders("my_custom_loader")defmy_custom_loader():
returnload_some_data()
assert"my_custom_loader"inml_datasets.loaders