refactored for code reuse

megatron/data_utils/corpora.py

-# 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,
 }