Commit cbd8c054 authored by Mohammad Shoeybi's avatar Mohammad Shoeybi
Browse files

refactored for code reuse

parent 34be7dd3

megatron/data_utils/corpora.py

+58 −106
Original line number Diff line number Diff line 
# coding=utf-8

# Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,8 +11,10 @@ 
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.



"""several datasets with preset arguments"""

from .datasets import json_dataset, csv_dataset

import os



class wikipedia(json_dataset):

	"""
@@ -21,11 +22,10 @@ class wikipedia(json_dataset): 


	command line usage: `--train-data wikipedia`

	"""

	#PATH = '/home/universal-lm-data.cosmos549/datasets/wikipedia/wikidump_lines.json'

	PATH = '/raid/mshoeybi/data/bert/wikipedia/wikidump_lines.json'

	assert_str = "make sure to set PATH at line 27 of data_utils/corpora.py"

	PATH = 'data/wikipedia/wikidump_lines.json'

	assert_str = "make sure to set PATH for wikipedia data_utils/corpora.py"

	def __init__(self, **kwargs):

		assert wikipedia.PATH != '<wikipedia_path>', \

		assert os.path.exists(wikipedia.PATH), \

                        wikipedia.assert_str

		if not kwargs:

			kwargs = {}
@@ -33,74 +33,26 @@ class wikipedia(json_dataset): 
		kwargs['loose_json'] = True

		super(wikipedia, self).__init__(wikipedia.PATH, **kwargs)



class roberta(json_dataset):

	"""

	dataset for roberta with arguments configured for convenience



	command line usage: `--train-data roberta`

class webtext(json_dataset):

	"""

	PATH = '/raid/mshoeybi/data/roberta/all_merged/rn_owt_sto_wiki_0.7_aug22/rn_owt_sto_wiki_0.7.json'

	assert_str = "make sure to set PATH at line 27 of data_utils/corpora.py"

	def __init__(self, **kwargs):

		assert roberta.PATH != '<roberta_path>', \

                                         roberta.assert_str

		if not kwargs:

			kwargs = {}

		kwargs['text_key'] = 'text'

		kwargs['loose_json'] = True

		super(roberta, self).__init__(roberta.PATH, **kwargs)

	dataset for webtext with arguments configured for convenience



class BooksCorpus(json_dataset):

        #PATH = '/home/universal-lm-data.cosmos549/datasets/BooksCorpus/books_lines.jsonl'

        PATH = '/raid/mshoeybi/data/bert/BooksCorpus/books_lines.jsonl'

        def __init__(self, **kwargs):

                if not kwargs:

                        kwargs = {}

                kwargs['text_key'] = 'text'

                kwargs['label_key'] = 'path'

                kwargs['loose_json'] = True

                super(BooksCorpus, self).__init__(BooksCorpus.PATH, **kwargs)



class Reddit(json_dataset):

        PATH = '/raid/mshoeybi/data/gpt2/adlr/urls_55M_ftNA_17M_sub_100_115_ftfy.json'

        #PATH='/home/universal-lm-data.cosmos549/datasets/OpenWebText/json_data/urls_55M_ftNA_17M_sub_100_115_ftfy.json'

        #PATH = '/raid/mshoeybi/data/gpt2/skylion007/openwebtext.jsonl'

        def __init__(self, **kwargs):

                if not kwargs:

                        kwargs = {}

                kwargs['text_key'] = 'text'

                kwargs['loose_json'] = True

                super(Reddit, self).__init__(Reddit.PATH, **kwargs)





class RedditAll(json_dataset):

        PATH = '/home/universal-lm-data.cosmos549/datasets/OpenWebText/json_data/reddit_all_ftfy.json'

        #PATH = '/raid/mshoeybi/data/gpt2/skylion007/openwebtext.jsonl'

        def __init__(self, **kwargs):

                if not kwargs:

                        kwargs = {}

                kwargs['text_key'] = 'text'

                kwargs['loose_json'] = True

                super(RedditAll, self).__init__(RedditAll.PATH, **kwargs)





class RedditAllLg200(json_dataset):

        PATH = '/home/universal-lm-data.cosmos549/datasets/OpenWebText/json_data/reddit_all_ftfy_lg200.json'

        #PATH = '/raid/mshoeybi/data/gpt2/skylion007/openwebtext.jsonl'

	command line usage: `--train-data webtext`

	"""

	PATH = 'data/webtext/data.json'

	assert_str = "make sure to set PATH for webtext data_utils/corpora.py"

	def __init__(self, **kwargs):

		assert os.path.exists(webtext.PATH), \

                        webtext.assert_str

		if not kwargs:

			kwargs = {}

		kwargs['text_key'] = 'text'

		kwargs['loose_json'] = True

                super(RedditAllLg200, self).__init__(RedditAllLg200.PATH, **kwargs)



		super(webtext, self).__init__(webtext.PATH, **kwargs)





NAMED_CORPORA = {

	'wikipedia': wikipedia,

        'roberta': roberta,

        'BooksCorpus': BooksCorpus,

        'Reddit': Reddit,

        'RedditAll': RedditAll,

        'RedditAllLg200': RedditAllLg200,

        'webtext': webtext,

}