pile.py 4.2 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
"""
The Pile: An 800GB Dataset of Diverse Text for Language Modeling
https://arxiv.org/pdf/2101.00027.pdf

The Pile is a 825 GiB diverse, open source language modelling data set that consists
of 22 smaller, high-quality datasets combined together. To score well on Pile
BPB (bits per byte), a model must be able to understand many disparate domains
including books, github repositories, webpages, chat logs, and medical, physics,
math, computer science, and philosophy papers.

Homepage: https://pile.eleuther.ai/
"""
Jason Phang's avatar
Jason Phang committed
13
14
15
16
17
18
19
20
21
22
23
import os

import lm_dataformat
import abc
import numpy as np
from lm_eval.base import rf, PerplexityTask
from ..metrics import mean, matthews_corrcoef, f1_score
from ..utils import general_detokenize
from best_download import download_file


24
25
26
27
28
29
30
31
32
33
_CITATION = """
@article{pile,
  title={The {P}ile: An 800GB Dataset of Diverse Text for Language Modeling},
  author={Gao, Leo and Biderman, Stella and Black, Sid and Golding, Laurence and Hoppe, Travis and Foster, Charles and Phang, Jason and He, Horace and Thite, Anish and Nabeshima, Noa and Presser, Shawn and Leahy, Connor},
  journal={arXiv preprint arXiv:2101.00027},
  year={2020}
}
"""


Leo Gao's avatar
Leo Gao committed
34
class PilePerplexityTask(PerplexityTask, abc.ABC):
35
    VERSION = 1
Jason Phang's avatar
Jason Phang committed
36
37
38
39
40
41

    PILE_SET_NAME = None
    VAL_PATH = 'data/pile/val.jsonl.zst'
    TEST_PATH = 'data/pile/test.jsonl.zst'

    def download(self):
Leo Gao's avatar
Leo Gao committed
42
        # TODO: separate pile val/test out by component so we don't have to scan the entire file once per set
Leo Gao's avatar
Leo Gao committed
43
        if not os.path.exists("data/pile/test.jsonl.zst"):
44
            # todo use new best_download fallback api
Leo Gao's avatar
Leo Gao committed
45
            os.makedirs("data/pile/", exist_ok=True)
46
47
            download_file("http://eaidata.bmk.sh/data/pile/val.jsonl.zst", local_file=self.VAL_PATH, expected_checksum="264c875d8bbd355d8daa9d032b75fd8fb91606218bb84dd1155b203fcd5fab92")
            download_file("http://eaidata.bmk.sh/data/pile/test.jsonl.zst", local_file=self.TEST_PATH, expected_checksum="0bb28c52d0b5596d389bf179ce2d43bf7f7ffae76b0d2d20b180c97f62e0975e")
Jason Phang's avatar
Jason Phang committed
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

    def validation_docs(self):
        rdr = lm_dataformat.Reader(self.VAL_PATH)
        for doc, metadata in rdr.stream_data(get_meta=True):
            if metadata["pile_set_name"] == self.PILE_SET_NAME:
                yield doc

    def test_docs(self):
        rdr = lm_dataformat.Reader(self.TEST_PATH)
        for doc, metadata in rdr.stream_data(get_meta=True):
            if metadata["pile_set_name"] == self.PILE_SET_NAME:
                yield doc

    def has_validation_docs(self):
        return True

    def has_test_docs(self):
        return True


Leo Gao's avatar
Leo Gao committed
68
class PileArxiv(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
69
70
71
    PILE_SET_NAME = "ArXiv"


Leo Gao's avatar
Leo Gao committed
72
class PileBooks3(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
73
74
75
    PILE_SET_NAME = "Books3"


Leo Gao's avatar
Leo Gao committed
76
class PileBookCorpus2(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
77
78
79
    PILE_SET_NAME = "BookCorpus2"


Leo Gao's avatar
Leo Gao committed
80
class PileDmMathematics(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
81
82
83
    PILE_SET_NAME = "DM Mathematics"


Leo Gao's avatar
Leo Gao committed
84
class PileEnron(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
85
86
87
    PILE_SET_NAME = "Enron Emails"


Leo Gao's avatar
Leo Gao committed
88
class PileEuroparl(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
89
90
91
    PILE_SET_NAME = "EuroParl"


Leo Gao's avatar
Leo Gao committed
92
class PileFreeLaw(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
93
94
95
    PILE_SET_NAME = "FreeLaw"


Leo Gao's avatar
Leo Gao committed
96
class PileGithub(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
97
98
99
    PILE_SET_NAME = "Github"


Leo Gao's avatar
Leo Gao committed
100
class PileGutenberg(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
101
102
103
    PILE_SET_NAME = "Gutenberg (PG-19)"


Leo Gao's avatar
Leo Gao committed
104
class PileHackernews(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
105
106
107
    PILE_SET_NAME = "HackerNews"


Leo Gao's avatar
Leo Gao committed
108
class PileNIHExporter(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
109
110
111
    PILE_SET_NAME = "NIH ExPorter"


Leo Gao's avatar
Leo Gao committed
112
class PileOpenSubtitles(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
113
114
115
    PILE_SET_NAME = "OpenSubtitles"


Leo Gao's avatar
Leo Gao committed
116
class PileOpenWebText2(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
117
118
119
    PILE_SET_NAME = "OpenWebText2"


Leo Gao's avatar
Leo Gao committed
120
class PilePhilPapers(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
121
122
123
    PILE_SET_NAME = "PhilPapers"


Leo Gao's avatar
Leo Gao committed
124
class PilePileCc(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
125
126
127
    PILE_SET_NAME = "Pile-CC"


Leo Gao's avatar
Leo Gao committed
128
class PilePubmedAbstracts(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
129
130
131
    PILE_SET_NAME = "PubMed Abstracts"


Leo Gao's avatar
Leo Gao committed
132
class PilePubmedCentral(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
133
134
135
    PILE_SET_NAME = "PubMed Central"


Leo Gao's avatar
Leo Gao committed
136
class PileStackExchange(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
137
138
139
    PILE_SET_NAME = "StackExchange"


Leo Gao's avatar
Leo Gao committed
140
class PileUspto(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
141
142
143
    PILE_SET_NAME = "USPTO Backgrounds"


Leo Gao's avatar
Leo Gao committed
144
class PileUbuntuIrc(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
145
    PILE_SET_NAME = "Ubuntu IRC"
Jason Phang's avatar
Jason Phang committed
146
147


Leo Gao's avatar
Leo Gao committed
148
class PileWikipedia(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
149
150
151
    PILE_SET_NAME = "Wikipedia (en)"


Leo Gao's avatar
Leo Gao committed
152
class PileYoutubeSubtitles(PilePerplexityTask):
Jason Phang's avatar
Jason Phang committed
153
    PILE_SET_NAME = "YoutubeSubtitles"