Skip to content

🌊 Machine learning dataset loaders for testing and example scripts

License

Notifications You must be signed in to change notification settings

explosion/ml-datasets

Folders and files

NameName
Last commit message
Last commit date

Latest commit

Β 

History

44 Commits
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

Machine learning dataset loaders for testing and examples

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

PyPi Version

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 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()

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()
for text, annot in train_data[0:5]:
    print(f"Review: {text}")
    print(f"Sentiment: {annot}")
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()
for text, annot in train_data[0:5]:
    print(f"Text: {text}")
    print(f"Category: {annot}")
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()
for text, annot in train_data[0:5]:
    print(f"Text: {text}")
    print(f"Genres: {annot}")
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()
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.

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")
def my_custom_loader():
    return load_some_data()

assert "my_custom_loader" in ml_datasets.loaders