update

liuzh47 · Jun 30, 2020 · fa011aa · fa011aa
1 parent 402d625
commit fa011aa
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/scripts/pretraining/pretraining_utils.py b/scripts/pretraining/pretraining_utils.py
@@ -262,11 +262,11 @@ def prepare_pretrain_text_dataset(
     """Create dataset based on the raw text files"""
     if not isinstance(filenames, (list, tuple)):
         filenames = [filenames]
-    # generate a filename based on the input filename ensuring no crash.
-    # filename example: urlsf_subset00-130_data.txt
-    suffix = re.findall(r'\d+-\d+', filenames[0])[0]
     if cached_file_path:
-        output_file = os.path.join(cached_file_path, "owt-pretrain-record-{}.npz".format(suffix))
+        # generate a filename based on the input filename ensuring no crash.
+        # filename example: urlsf_subset00-130_data.txt
+        suffix = re.split(r'\.|/', filenames[0])[-2]
+        output_file = os.path.join(cached_file_path, "{}-pretrain-record.npz".format(suffix))
     else:
         output_file = None
     np_features = get_all_features(

diff --git a/scripts/pretraining/run_electra.py b/scripts/pretraining/run_electra.py
@@ -273,7 +273,7 @@ def train(args):
                  'num_workers: {}, rank: {}'.format(
                      args.num_buckets, num_workers, rank))
     if args.from_raw_text:
-        if not os.path.exists(args.cached_file_path):
+        if args.cached_file_path and not os.path.exists(args.cached_file_path):
             os.mkdir(args.cached_file_path)
         get_dataset_fn = functools.partial(get_pretrain_data_text,
                                            max_seq_length=args.max_seq_length,