Loaders for various machine learning datasets for testing and example scripts. Previously in thinc.extra.datasets
.
The package can be installed via pip:
pip install ml-datasets
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 directly from ml_datasets import imdb train_data, dev_data = imdb()
# Load via registry from ml_datasets import loaders imdb_loader = loaders.get("imdb") train_data, dev_data = imdb_loader()
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 |
ID / Function | Description | ML task | From URL |
---|---|---|---|
mnist | MNIST data | Image recognition | β |
Each instance contains the text of a movie review, and a sentiment expressed as 0
or 1
.
train_data, dev_data = ml_datasets.imdb() for text, annot in train_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) |
Each instance contains an ontological description, and a classification into one of the 14 distinct labels.
train_data, dev_data = ml_datasets.dbpedia() for text, annot in train_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 |
Each instance contains a movie description, and a classification into a list of appropriate genres.
train_data, dev_data = ml_datasets.cmu() for text, annot in train_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 | - |
train_data, dev_data = ml_datasets.quora_questions() for questions, annot in train_data[0:50]: q1, q2 = questions print(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 |
Loaders can be registered externally using the loaders
registry as a decorator. For example:
@ml_datasets.loaders("my_custom_loader") def my_custom_loader(): return load_some_data() assert "my_custom_loader" in ml_datasets.loaders