feat: remove mpi4py

src/improved_diffusion/dist_util.py

@@ -8,7 +8,6 @@ import socket
 
 import blobfile as bf
 
-from mpi4py import MPI
 import torch as th
 import torch.distributed as dist
 
@@ -46,26 +45,26 @@ def setup_dist(rank, world_size, port="12145"):
     dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
 
 
-def dev():
-    """
-    Get the device to use for torch.distributed.
-    """
-    if th.cuda.is_available():
-        return th.device(f"cuda:{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}")
-    return th.device("cpu")
-
-
-def load_state_dict(path, **kwargs):
-    """
-    Load a PyTorch file without redundant fetches across MPI ranks.
-    """
-    if MPI.COMM_WORLD.Get_rank() == 0:
-        with bf.BlobFile(path, "rb") as f:
-            data = f.read()
-    else:
-        data = None
-    data = MPI.COMM_WORLD.bcast(data)
-    return th.load(io.BytesIO(data), **kwargs)
+# def dev():
+#     """
+#     Get the device to use for torch.distributed.
+#     """
+#     if th.cuda.is_available():
+#         return th.device(f"cuda:{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}")
+#     return th.device("cpu")
+
+
+# def load_state_dict(path, **kwargs):
+#     """
+#     Load a PyTorch file without redundant fetches across MPI ranks.
+#     """
+#     if MPI.COMM_WORLD.Get_rank() == 0:
+#         with bf.BlobFile(path, "rb") as f:
+#             data = f.read()
+#     else:
+#         data = None
+#     data = MPI.COMM_WORLD.bcast(data)
+#     return th.load(io.BytesIO(data), **kwargs)
 
 
 def sync_params(params):

src/improved_diffusion/text_datasets.py

@@ -1,13 +1,21 @@
 # from PIL import Image
 # import blobfile as bf
-from mpi4py import MPI
+# from mpi4py import MPI
 import numpy as np
 from torch.utils.data import DataLoader, Dataset
-from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer, default_data_collator, PreTrainedTokenizerFast, \
-    PreTrainedTokenizer
+from transformers import (
+    AutoModelForCausalLM,
+    AutoConfig,
+    AutoTokenizer,
+    default_data_collator,
+    PreTrainedTokenizerFast,
+    PreTrainedTokenizer,
+)
+
 # from datasets import load_dataset
 import sys, os
 import torch
+
 # sys.path.insert(0, os.path.join(sys.path[0], '../../transformers/examples/pytorch/language-modeling'))
 # from custom_trainer import GPT2LMHeadModelCompress, BERTModelCompress, AutoEncoderWithNoise
 from collections import Counter, defaultdict
@@ -16,8 +24,18 @@ from itertools import chain
 
 
 def load_data_text(
-    *, data_dir, batch_size, image_size, class_cond=False, deterministic=False, data_args=None, 
-        task_mode='roc', model=None, padding_mode='block', split='train', load_vocab=None,
+    *,
+    data_dir,
+    batch_size,
+    image_size,
+    class_cond=False,
+    deterministic=False,
+    data_args=None,
+    task_mode="roc",
+    model=None,
+    padding_mode="block",
+    split="train",
+    load_vocab=None,
 ):
     """
     For a dataset, create a generator over (images, kwargs) pairs.
@@ -35,29 +53,34 @@ def load_data_text(
                        exception will be raised.
     :param deterministic: if True, yield results in a deterministic order.
     """
-    print('hello loading text data. ')
+    print("hello loading text data. ")
 
-    if data_args.experiment.startswith('random') and model is None:
+    if data_args.experiment.startswith("random") and model is None:
         model = None
     # elif data_args.experiment.startswith('random') and model is not None:
     #     print('loading initialized random embeddings. ')
 
-    if task_mode == 'roc' or task_mode == 'roc-aug' :
+    if task_mode == "roc" or task_mode == "roc-aug":
         pass
         # training_data, model = get_corpus_rocstory(data_args, model, image_size,
         #                                     padding_mode=padding_mode, split=split,
-                                            # load_vocab=load_vocab)
-    elif task_mode == 'simple-wiki':
+        # load_vocab=load_vocab)
+    elif task_mode == "simple-wiki":
         pass
         # training_data, model = get_corpus_rocstory(data_args, model, image_size,
-                                            # padding_mode=padding_mode, split=split,
-                                            # load_vocab=load_vocab)
-
-    elif task_mode == 'e2e-tgt':
-        print('hello loading e2e-tgt. ')
-        training_data, model = get_corpus_rocstory(data_args, model, image_size,
-                                            padding_mode=padding_mode, split=split,
-                                            load_vocab=load_vocab)
+        # padding_mode=padding_mode, split=split,
+        # load_vocab=load_vocab)
+
+    elif task_mode == "e2e-tgt":
+        print("hello loading e2e-tgt. ")
+        training_data, model = get_corpus_rocstory(
+            data_args,
+            model,
+            image_size,
+            padding_mode=padding_mode,
+            split=split,
+            load_vocab=load_vocab,
+        )
     # elif task_mode == 'yelp':
     #     print('hello loading yelp ')
     #     training_data, model = get_corpus_rocstory(data_args, model, image_size,
@@ -80,8 +103,12 @@ def load_data_text(
     #     training_data, model = get_corpus_book(data_args, tokenizer, model, image_size,
     #                                           padding_mode=padding_mode, split=split,)
 
-    if data_args.modality in ['roc-aug', 'roc', 'book', 'yelp', 'commonGen', 'commonGen-aug'] and data_args.cache_mode=='no':
-        pass# dataset = TextDataset_NoCache(
+    if (
+        data_args.modality
+        in ["roc-aug", "roc", "book", "yelp", "commonGen", "commonGen-aug"]
+        and data_args.cache_mode == "no"
+    ):
+        pass  # dataset = TextDataset_NoCache(
         #     training_data,
         #     image_size,
         #     data_args,
@@ -98,7 +125,7 @@ def load_data_text(
 
     if deterministic:
 
-        pass# data_loader = DataLoader(
+        pass  # data_loader = DataLoader(
         #     dataset,
         #     batch_size=batch_size,  # 20,
         #     drop_last=True,
@@ -117,64 +144,83 @@ def load_data_text(
     while True:
         yield from data_loader
 
+
 def helper_tokenize_encode_cond(sentence_lst, vocab_dict, model, seqlen, data_args):
     result_train_lst = []
     group_lst = defaultdict(list)
     with torch.no_grad():
-        for (src_ids, input_ids) in sentence_lst:
-            tokenized_ = [vocab_dict.get(x, vocab_dict['UNK']) for x in input_ids]
-            tokenized_src = [vocab_dict.get(x, vocab_dict['UNK']) for x in src_ids]
+        for src_ids, input_ids in sentence_lst:
+            tokenized_ = [vocab_dict.get(x, vocab_dict["UNK"]) for x in input_ids]
+            tokenized_src = [vocab_dict.get(x, vocab_dict["UNK"]) for x in src_ids]
             input_ids = [0] + tokenized_ + [1]
-            group_lst['word_ids'].append(input_ids)
-            group_lst['src_ids'].append(tokenized_src)
+            group_lst["word_ids"].append(input_ids)
+            group_lst["src_ids"].append(tokenized_src)
 
-        print(group_lst['word_ids'][:2])
-        print('padding mode is pad')
+        print(group_lst["word_ids"][:2])
+        print("padding mode is pad")
         max_length = seqlen
-        group_lst['word_ids'] = _collate_batch_helper(group_lst['word_ids'], vocab_dict['PAD'], max_length)
-        max_src_length = max([len(xx) for xx in group_lst['src_ids']])
+        group_lst["word_ids"] = _collate_batch_helper(
+            group_lst["word_ids"], vocab_dict["PAD"], max_length
+        )
+        max_src_length = max([len(xx) for xx in group_lst["src_ids"]])
         print(max_src_length, seqlen)
         max_src_length = min(seqlen, max_src_length)
-        group_lst['src_ids'], group_lst['src_mask'] = _collate_batch_helper(group_lst['src_ids'],
-                                                                            vocab_dict['PAD'],
-                                                                            max_src_length,
-                                                                            return_mask=True)
-
+        group_lst["src_ids"], group_lst["src_mask"] = _collate_batch_helper(
+            group_lst["src_ids"], vocab_dict["PAD"], max_src_length, return_mask=True
+        )
 
-        for input_ids, src_ids, src_mask in zip(group_lst['word_ids'], group_lst['src_ids'],
-                                      group_lst['src_mask']):
-            if data_args.experiment.startswith('random'):
+        for input_ids, src_ids, src_mask in zip(
+            group_lst["word_ids"], group_lst["src_ids"], group_lst["src_mask"]
+        ):
+            if data_args.experiment.startswith("random"):
                 hidden_state = model(torch.tensor(input_ids))
-            elif data_args.experiment == 'gpt2_pre_compress':
+            elif data_args.experiment == "gpt2_pre_compress":
                 input_ids2 = torch.tensor(input_ids).to(model.device)
                 input_embs = model.transformer.wte(input_ids2)  # input_embs
                 hidden_state = model.down_proj(input_embs)
                 hidden_state = hidden_state * data_args.emb_scale_factor
-            result_train_lst.append({'input_ids': input_ids,
-                                     'hidden_states': hidden_state.cpu().tolist(),
-                                     'src_ids':src_ids,
-                                     'src_mask':src_mask
-                                     })
+            result_train_lst.append(
+                {
+                    "input_ids": input_ids,
+                    "hidden_states": hidden_state.cpu().tolist(),
+                    "src_ids": src_ids,
+                    "src_mask": src_mask,
+                }
+            )
 
     return result_train_lst
 
-def helper_tokenize_stream(sentence_lst, vocab_dict, model, seqlen, data_args, padding_mode, ):
+
+def helper_tokenize_stream(
+    sentence_lst,
+    vocab_dict,
+    model,
+    seqlen,
+    data_args,
+    padding_mode,
+):
     import psutil
+
     # Process.memory_info is expressed in bytes, so convert to megabytes
     print(f"RAM used: {psutil.Process().memory_info().rss / (1024 * 1024):.2f} MB")
     from datasets import Dataset as Dataset2
-    raw_datasets = Dataset2.from_dict({'text':sentence_lst})
+
+    raw_datasets = Dataset2.from_dict({"text": sentence_lst})
     print(raw_datasets)
     print(f"RAM used: {psutil.Process().memory_info().rss / (1024 * 1024):.2f} MB")
 
-
     def tokenize_function(examples):
         if isinstance(vocab_dict, dict):
-            input_ids = [[0] + [vocab_dict.get(x, vocab_dict['UNK']) for x in seq] + [1] for seq in examples['text']]
+            input_ids = [
+                [0] + [vocab_dict.get(x, vocab_dict["UNK"]) for x in seq] + [1]
+                for seq in examples["text"]
+            ]
         elif isinstance(vocab_dict, PreTrainedTokenizerFast):
-            examples['text'] = [" ".join(seq) for seq in examples['text']]
-            input_ids = vocab_dict(examples['text'], add_special_tokens=True)['input_ids']
-        result_dict = {'input_ids': input_ids}
+            examples["text"] = [" ".join(seq) for seq in examples["text"]]
+            input_ids = vocab_dict(examples["text"], add_special_tokens=True)[
+                "input_ids"
+            ]
+        result_dict = {"input_ids": input_ids}
         # clm input could be much much longer than block_size
         return result_dict
 
@@ -182,28 +228,30 @@ def helper_tokenize_stream(sentence_lst, vocab_dict, model, seqlen, data_args, p
         tokenize_function,
         batched=True,
         num_proc=4,
-        remove_columns=['text'],
+        remove_columns=["text"],
         load_from_cache_file=True,
         desc="Running tokenizer on dataset",
     )
     print(tokenized_datasets)
     print(f"RAM used: {psutil.Process().memory_info().rss / (1024 * 1024):.2f} MB")
 
-    if padding_mode == 'block':
+    if padding_mode == "block":
         block_size = seqlen
+
         def group_texts(examples):
-            concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
+            concatenated_examples = {
+                k: list(chain(*examples[k])) for k in examples.keys()
+            }
             total_length = len(concatenated_examples[list(examples.keys())[0]])
             if total_length >= block_size:
                 total_length = (total_length // block_size) * block_size
             result = {
-                k: [t[i: i + block_size] for i in range(0, total_length, block_size)]
+                k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
                 for k, t in concatenated_examples.items()
             }
             result["labels"] = result["input_ids"].copy()
             return result
 
-
         lm_datasets = tokenized_datasets.map(
             group_texts,
             batched=True,
@@ -212,12 +260,17 @@ def helper_tokenize_stream(sentence_lst, vocab_dict, model, seqlen, data_args, p
             desc=f"Grouping texts in chunks of {block_size}",
         )
     else:
+
         def pad_function(group_lst):
             max_length = seqlen
             if isinstance(vocab_dict, dict):
-                group_lst['input_ids'] = _collate_batch_helper(group_lst['input_ids'], vocab_dict['PAD'], max_length)
+                group_lst["input_ids"] = _collate_batch_helper(
+                    group_lst["input_ids"], vocab_dict["PAD"], max_length
+                )
             else:
-                group_lst['input_ids'] = _collate_batch_helper(group_lst['input_ids'], vocab_dict.pad_token_id, max_length)
+                group_lst["input_ids"] = _collate_batch_helper(
+                    group_lst["input_ids"], vocab_dict.pad_token_id, max_length
+                )
             return group_lst
 
         # Process.memory_info is expressed in bytes, so convert to megabytes
@@ -230,59 +283,72 @@ def helper_tokenize_stream(sentence_lst, vocab_dict, model, seqlen, data_args, p
             desc=f"padding",
         )
 
-
-    print(lm_datasets, 'padded dataset')
+    print(lm_datasets, "padded dataset")
     print(f"RAM used: {psutil.Process().memory_info().rss / (1024 * 1024):.2f} MB")
     import datasets
+
     raw_datasets = datasets.DatasetDict()
-    raw_datasets['train'] = lm_datasets
+    raw_datasets["train"] = lm_datasets
     print(f"RAM used: {psutil.Process().memory_info().rss / (1024 * 1024):.2f} MB")
     return raw_datasets
 
-def helper_tokenize_encode(sentence_lst, vocab_dict, model, seqlen, data_args, padding_mode, ):
+
+def helper_tokenize_encode(
+    sentence_lst,
+    vocab_dict,
+    model,
+    seqlen,
+    data_args,
+    padding_mode,
+):
     result_train_lst = []
     group_lst = defaultdict(list)
     with torch.no_grad():
         for input_ids in sentence_lst:
-            tokenized_ = [vocab_dict.get(x, vocab_dict['UNK']) for x in input_ids]
+            tokenized_ = [vocab_dict.get(x, vocab_dict["UNK"]) for x in input_ids]
             input_ids = [0] + tokenized_ + [1]
-            group_lst['word_ids'].append(input_ids)
-        print(group_lst['word_ids'][:2])
+            group_lst["word_ids"].append(input_ids)
+        print(group_lst["word_ids"][:2])
 
-        if padding_mode == 'block':
-            print('padding mode is block')
+        if padding_mode == "block":
+            print("padding mode is block")
             concatenated_examples = {k: sum(group_lst[k], []) for k in group_lst.keys()}
             total_length = len(concatenated_examples[list(group_lst.keys())[0]])
             block_size = seqlen
             total_length = (total_length // block_size) * block_size
             # Split by chunks of max_len.
             group_lst = {
-                k: [t[i: i + block_size] for i in range(0, total_length, block_size)]
+                k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
                 for k, t in concatenated_examples.items()
             }
-        elif padding_mode == 'pad':
-            print('padding mode is pad')
+        elif padding_mode == "pad":
+            print("padding mode is pad")
             max_length = seqlen
-            group_lst['word_ids'] = _collate_batch_helper(group_lst['word_ids'], vocab_dict['PAD'], max_length)
+            group_lst["word_ids"] = _collate_batch_helper(
+                group_lst["word_ids"], vocab_dict["PAD"], max_length
+            )
 
-        for input_ids in group_lst['word_ids']:
-            if data_args.experiment.startswith('random'):
+        for input_ids in group_lst["word_ids"]:
+            if data_args.experiment.startswith("random"):
                 hidden_state = model(torch.tensor(input_ids))
-            elif data_args.experiment == 'gpt2_pre_compress':
+            elif data_args.experiment == "gpt2_pre_compress":
                 input_ids2 = torch.tensor(input_ids).to(model.device)
                 input_embs = model.transformer.wte(input_ids2)  # input_embs
                 hidden_state = model.down_proj(input_embs)
                 hidden_state = hidden_state * data_args.emb_scale_factor
-            elif data_args.experiment == 'glove':
+            elif data_args.experiment == "glove":
                 hidden_state = model(torch.tensor(input_ids))
-            result_train_lst.append({'input_ids': input_ids, 'hidden_states': hidden_state.cpu().tolist()})
+            result_train_lst.append(
+                {"input_ids": input_ids, "hidden_states": hidden_state.cpu().tolist()}
+            )
 
     return result_train_lst
 
+
 def load_glove_model(File):
     print("Loading Glove Model")
     glove_model = {}
-    with open(File,'r') as f:
+    with open(File, "r") as f:
         for line in f:
             split_line = line.split()
             word = split_line[0]
@@ -292,9 +358,10 @@ def load_glove_model(File):
     print(f"{len(glove_model)} words loaded!")
     return glove_model
 
+
 def load_glove(vocab):
     model = torch.nn.Embedding(len(vocab), 50)
-    glove_model = load_glove_model('predictability/glove/glove.6B.50d.txt')
+    glove_model = load_glove_model("predictability/glove/glove.6B.50d.txt")
     array_lst = []
     count_ = 0
     for word, idx in vocab.items():
@@ -303,20 +370,21 @@ def load_glove(vocab):
         else:
             count_ += 1
             array_lst.append(torch.randn(50))
-    print(f'{count_} out of {len(vocab)} is initialized. ')
+    print(f"{count_} out of {len(vocab)} is initialized. ")
     array_lst = torch.stack(array_lst)
     print(torch.norm(array_lst, dim=-1).mean())
     model.weight.data = array_lst
     return model
 
 
-def get_corpus_rocstory(data_args, model, image_size, padding_mode='block',
-                        split='train', load_vocab=None):
+def get_corpus_rocstory(
+    data_args, model, image_size, padding_mode="block", split="train", load_vocab=None
+):
     import csv, torch, json
     from spacy.lang.en import English
 
-    if data_args.experiment_mode == 'lm':
-        if data_args.modality == 'roc':
+    if data_args.experiment_mode == "lm":
+        if data_args.modality == "roc":
             pass
             # print('loading dataset from ROCStory')
             # nlp = English()
@@ -347,7 +415,7 @@ def get_corpus_rocstory(data_args, model, image_size, padding_mode='block',
             # #         sentence_lst.append(word_lst)
             # # sentence_lst = sentence_lst[1:]
             # print(sentence_lst[:2])
-        if data_args.modality == 'roc-aug':
+        if data_args.modality == "roc-aug":
             pass
             # print('loading dataset from ROCStory')
             # nlp = English()
@@ -381,7 +449,7 @@ def get_corpus_rocstory(data_args, model, image_size, padding_mode='block',
             #                 word_lst = [x.text for x in tokenizer(sentences)]
             #                 sentence_lst.append(word_lst)
             # print(sentence_lst[:2],sentence_lst[-2:], 'dataset size=',len(sentence_lst))
-        elif data_args.modality == 'simple-wiki':
+        elif data_args.modality == "simple-wiki":
             pass
             # print('loading dataset from simple wikipedia')
             # sentence_lst = []
@@ -390,57 +458,62 @@ def get_corpus_rocstory(data_args, model, image_size, padding_mode='block',
             #         word_lst = row.lower().split()
             #         sentence_lst.append(word_lst)
             # print(sentence_lst[:2])
-        elif data_args.modality == 'e2e-tgt':
-            print('loading dataset from simple e2e dataset')
+        elif data_args.modality == "e2e-tgt":
+            print("loading dataset from simple e2e dataset")
             sentence_lst = []
             nlp = English()
             tokenizer = nlp.tokenizer
-            if split == 'train':
-                print('loading form the TRAIN set')
-                path = '/data0/gonghaisong/Diffusion-LM/datasets/e2e_data/src1_train.txt'
+            if split == "train":
+                print("loading form the TRAIN set")
+                path = (
+                    "/data0/gonghaisong/Diffusion-LM/datasets/e2e_data/src1_train.txt"
+                )
                 # path = f'../{data_args.e2e_train}/src1_train.txt'
-            elif split == 'valid':
-                print('loading form the VALID set')
-                path = f'../{data_args.e2e_train}/src1_valid.txt'
-                path = '/data0/gonghaisong/Diffusion-LM/datasets/e2e_data/src1_valid.txt'
-            elif split == 'test':
-                print('loading form the TEST set')
-                path = f'../{data_args.e2e_train}/src1_test.txt'
-                path = '/data0/gonghaisong/Diffusion-LM/datasets/e2e_data/src1_test.txt'
-            elif split == 'debug':
-                print('loading form the DEBUG set')
+            elif split == "valid":
+                print("loading form the VALID set")
+                path = f"../{data_args.e2e_train}/src1_valid.txt"
+                path = (
+                    "/data0/gonghaisong/Diffusion-LM/datasets/e2e_data/src1_valid.txt"
+                )
+            elif split == "test":
+                print("loading form the TEST set")
+                path = f"../{data_args.e2e_train}/src1_test.txt"
+                path = "/data0/gonghaisong/Diffusion-LM/datasets/e2e_data/src1_test.txt"
+            elif split == "debug":
+                print("loading form the DEBUG set")
                 path = data_args.debug_path
                 import json
-                with open(path, 'r') as ff:
+
+                with open(path, "r") as ff:
                     for line in ff:
-                        sentence_lst.append(json.loads(line)[0].split(' '))
+                        sentence_lst.append(json.loads(line)[0].split(" "))
                 sentence_lst = sentence_lst + sentence_lst
-            if split in ['train', 'valid', 'test']:
-                with open(path, 'r') as ff:
+            if split in ["train", "valid", "test"]:
+                with open(path, "r") as ff:
                     for row in ff:
-                        word_lst = row.split('||')[1]
+                        word_lst = row.split("||")[1]
                         word_lst = [x.text for x in tokenizer(word_lst)]
                         sentence_lst.append(word_lst)
             print(sentence_lst[:2])
 
-        elif data_args.modality == 'yelp':
-            print('loading dataset from simple YelpNLG dataset')
+        elif data_args.modality == "yelp":
+            print("loading dataset from simple YelpNLG dataset")
             sentence_lst = []
             nlp = English()
             tokenizer = nlp.tokenizer
-            if split == 'train':
-                print('loading form the TRAIN set')
-                path = f'{data_args.yelp_train}/yelpnlg-train.csv'
-            elif split == 'valid':
-                print('loading form the VALID set')
-                path = f'{data_args.yelp_train}/yelpnlg-dev.csv'
-            elif split == 'test':
-                print('loading form the TEST set')
-                path = f'{data_args.yelp_train}/yelpnlg-test.csv'
-            if split in ['train', 'valid', 'test']:
-
-                with open(path, 'r') as csvfile:
-                    yelp_reader = csv.reader(csvfile) #delimiter=' ', quotechar='|')
+            if split == "train":
+                print("loading form the TRAIN set")
+                path = f"{data_args.yelp_train}/yelpnlg-train.csv"
+            elif split == "valid":
+                print("loading form the VALID set")
+                path = f"{data_args.yelp_train}/yelpnlg-dev.csv"
+            elif split == "test":
+                print("loading form the TEST set")
+                path = f"{data_args.yelp_train}/yelpnlg-test.csv"
+            if split in ["train", "valid", "test"]:
+
+                with open(path, "r") as csvfile:
+                    yelp_reader = csv.reader(csvfile)  # delimiter=' ', quotechar='|')
                     for row in yelp_reader:
                         sentences = row[1]
                         word_lst = [x.text for x in tokenizer(sentences)]
@@ -448,175 +521,188 @@ def get_corpus_rocstory(data_args, model, image_size, padding_mode='block',
                 sentence_lst = sentence_lst[1:]
             print(sentence_lst[:2])
 
-        elif data_args.modality == 'commonGen':
-            print('loading dataset from simple YelpNLG dataset')
+        elif data_args.modality == "commonGen":
+            print("loading dataset from simple YelpNLG dataset")
             sentence_lst = []
             nlp = English()
             tokenizer = nlp.tokenizer
-            if split == 'train':
-                print('loading form the TRAIN set')
-                path = f'{data_args.commonGen_train}/commongen.train.jsonl'
-            elif split == 'valid':
-                print('loading form the VALID set')
-                path = f'{data_args.commonGen_train}/commongen.dev.jsonl'
-            elif split == 'test':
-                print('loading form the TEST set')
-                path = f'{data_args.commonGen_train}/commongen.test.jsonl'
-            if split in ['train', 'valid', 'test']:
-                with open(path, 'r') as ff:
+            if split == "train":
+                print("loading form the TRAIN set")
+                path = f"{data_args.commonGen_train}/commongen.train.jsonl"
+            elif split == "valid":
+                print("loading form the VALID set")
+                path = f"{data_args.commonGen_train}/commongen.dev.jsonl"
+            elif split == "test":
+                print("loading form the TEST set")
+                path = f"{data_args.commonGen_train}/commongen.test.jsonl"
+            if split in ["train", "valid", "test"]:
+                with open(path, "r") as ff:
                     for line in ff:
                         line = json.loads(line)
-                        for sentences in line['scene']:
+                        for sentences in line["scene"]:
                             word_lst = [x.text for x in tokenizer(sentences)]
                             sentence_lst.append(word_lst)
             print(sentence_lst[:2])
 
-        elif data_args.modality == 'commonGen-aug':
-            print('loading dataset from simple YelpNLG dataset')
+        elif data_args.modality == "commonGen-aug":
+            print("loading dataset from simple YelpNLG dataset")
             sentence_lst = []
             nlp = English()
             tokenizer = nlp.tokenizer
-            if split == 'train':
-                print('loading form the TRAIN set')
-                path = f'{data_args.commonGen_train}/commongen.train.jsonl'
-                path_lst = [f'{data_args.roc_train}/roc_train.json']
-                path_lst.append('diffusion_lm/improved-diffusion/diff_models/rocstories_gptj.txt')
-            elif split == 'valid':
-                print('loading form the VALID set')
-                path = f'{data_args.commonGen_train}/commongen.dev.jsonl'
+            if split == "train":
+                print("loading form the TRAIN set")
+                path = f"{data_args.commonGen_train}/commongen.train.jsonl"
+                path_lst = [f"{data_args.roc_train}/roc_train.json"]
+                path_lst.append(
+                    "diffusion_lm/improved-diffusion/diff_models/rocstories_gptj.txt"
+                )
+            elif split == "valid":
+                print("loading form the VALID set")
+                path = f"{data_args.commonGen_train}/commongen.dev.jsonl"
                 path_lst = []
-            elif split == 'test':
-                print('loading form the TEST set')
-                path = f'{data_args.commonGen_train}/commongen.test.jsonl'
+            elif split == "test":
+                print("loading form the TEST set")
+                path = f"{data_args.commonGen_train}/commongen.test.jsonl"
                 path_lst = []
 
-            if split in ['train', 'valid', 'test']:
-                with open(path, 'r') as ff:
+            if split in ["train", "valid", "test"]:
+                with open(path, "r") as ff:
                     for line in ff:
                         line = json.loads(line)
-                        for sentences in line['scene']:
+                        for sentences in line["scene"]:
                             word_lst = [x.text for x in tokenizer(sentences)]
                             sentence_lst.append(word_lst)
             print(sentence_lst[:2])
             import itertools
+
             for path in path_lst:
-                if path.endswith('txt'):
-                    with open(path, 'r') as roc_reader:
+                if path.endswith("txt"):
+                    with open(path, "r") as roc_reader:
                         for row in roc_reader:
                             sentences = row.strip()
                             word_lst = [x.text for x in tokenizer(sentences)]
                             spl = [[]]
-                            for x, y in itertools.groupby(word_lst, lambda z: z == '.'):
+                            for x, y in itertools.groupby(word_lst, lambda z: z == "."):
                                 spl[-1].extend(y)
-                                if x: spl.append([])
+                                if x:
+                                    spl.append([])
                             sentence_lst.extend(spl[:-1])
                 else:
-                    with open(path, 'r') as roc_reader:
+                    with open(path, "r") as roc_reader:
                         for row in roc_reader:
                             sentences = json.loads(row)[0].strip()
                             word_lst = [x.text for x in tokenizer(sentences)]
                             spl = [[]]
-                            for x, y in itertools.groupby(word_lst, lambda z: z == '.'):
+                            for x, y in itertools.groupby(word_lst, lambda z: z == "."):
                                 spl[-1].extend(y)
-                                if x: spl.append([])
+                                if x:
+                                    spl.append([])
                             sentence_lst.extend(spl[:-1])
 
             print(sentence_lst[-2:])
 
-
         # get tokenizer.
         if load_vocab is None:
             counter = Counter()
             for input_ids in sentence_lst:
                 counter.update(input_ids)
 
-    if data_args.experiment_mode == 'conditional_gen':
-        if data_args.modality == 'e2e':
-            print('loading dataset from simple e2e dataset')
+    if data_args.experiment_mode == "conditional_gen":
+        if data_args.modality == "e2e":
+            print("loading dataset from simple e2e dataset")
             sentence_lst = []
             nlp = English()
             tokenizer = nlp.tokenizer
-            if split == 'train':
-                path = f'{data_args.e2e_train}/src1_train.txt'
-                with open(path, 'r') as ff:
+            if split == "train":
+                path = f"{data_args.e2e_train}/src1_train.txt"
+                with open(path, "r") as ff:
                     for row in ff:
-                        src_lst, word_lst = row.split('||')
+                        src_lst, word_lst = row.split("||")
                         word_lst = [x.text for x in tokenizer(word_lst)]
                         src_lst = [x.text for x in tokenizer(src_lst)]
                         sentence_lst.append((src_lst, word_lst))
-            elif split == 'valid':
-                path = f'{data_args.e2e_train}/src1_valid.txt'
+            elif split == "valid":
+                path = f"{data_args.e2e_train}/src1_valid.txt"
                 sentence_lst = read_e2e_files(path, data_args, tokenizer)
             print(sentence_lst[:2])
         # get tokenizer.
         if load_vocab is None:
             counter = Counter()
-            for (src_ids, input_ids) in sentence_lst:
+            for src_ids, input_ids in sentence_lst:
                 counter.update(input_ids)
                 counter.update(src_ids)
 
     if load_vocab is None:
-        vocab_dict = {'START': 0, 'END': 1, 'UNK':2, 'PAD':3}
+        vocab_dict = {"START": 0, "END": 1, "UNK": 2, "PAD": 3}
         for k, v in counter.items():
             if v > 10:
                 vocab_dict[k] = len(vocab_dict)
         print(len(counter), len(vocab_dict))
 
-        path_save_vocab = '/data0/gonghaisong/Diffusion-LM/improved-diffusion/diffusion_models/diff_e2e-tgt_block_rand16_transformer_lr0.0001_0.0_2000_sqrt_Lsimple_h128_s2_d0.1_sd102_xstart_e2e/vocab.json'
-        print(f'save the vocab to {path_save_vocab}')
-        with open(path_save_vocab, 'w') as f:
+        path_save_vocab = "/data0/gonghaisong/Diffusion-LM/improved-diffusion/diffusion_models/diff_e2e-tgt_block_rand16_transformer_lr0.0001_0.0_2000_sqrt_Lsimple_h128_s2_d0.1_sd102_xstart_e2e/vocab.json"
+        print(f"save the vocab to {path_save_vocab}")
+        with open(path_save_vocab, "w") as f:
             json.dump(vocab_dict, f)
     else:
         vocab_dict = load_vocab
-        path_save_vocab = '/data0/gonghaisong/Diffusion-LM/improved-diffusion/diffusion_models/diff_e2e-tgt_block_rand16_transformer_lr0.0001_0.0_2000_sqrt_Lsimple_h128_s2_d0.1_sd102_xstart_e2e/vocab.json'
+        path_save_vocab = "/data0/gonghaisong/Diffusion-LM/improved-diffusion/diffusion_models/diff_e2e-tgt_block_rand16_transformer_lr0.0001_0.0_2000_sqrt_Lsimple_h128_s2_d0.1_sd102_xstart_e2e/vocab.json"
         if not os.path.exists(path_save_vocab):
-            print(f'save the vocab to {path_save_vocab}')
+            print(f"save the vocab to {path_save_vocab}")
             if isinstance(vocab_dict, dict):
-                with open(path_save_vocab, 'w') as f:
+                with open(path_save_vocab, "w") as f:
                     json.dump(vocab_dict, f)
-                assert vocab_dict['START'] == 0
+                assert vocab_dict["START"] == 0
             elif isinstance(vocab_dict, PreTrainedTokenizerFast):
                 vocab_dict.save_pretrained(data_args.checkpoint_path)
             else:
                 assert False, "invalid type of vocab_dict"
 
-
-
-    if model is None and data_args.experiment == 'random':
+    if model is None and data_args.experiment == "random":
         model = torch.nn.Embedding(len(vocab_dict), data_args.in_channel)
-        print('initializing the random embeddings', model)
+        print("initializing the random embeddings", model)
         torch.nn.init.normal_(model.weight)
-        path_save = '/data0/gonghaisong/Diffusion-LM/improved-diffusion/diffusion_models/diff_e2e-tgt_block_rand16_transformer_lr0.0001_0.0_2000_sqrt_Lsimple_h128_s2_d0.1_sd102_xstart_e2e/random_emb.torch'
-        print(f'save the random encoder to {data_args.checkpoint_path}/random_emb.torch')
+        path_save = "/data0/gonghaisong/Diffusion-LM/improved-diffusion/diffusion_models/diff_e2e-tgt_block_rand16_transformer_lr0.0001_0.0_2000_sqrt_Lsimple_h128_s2_d0.1_sd102_xstart_e2e/random_emb.torch"
+        print(
+            f"save the random encoder to {data_args.checkpoint_path}/random_emb.torch"
+        )
         torch.save(model.state_dict(), path_save)
 
     # path_save = f'{data_args.checkpoint_path}/random_emb.torch'
     # if not os.path.exists(path_save) and data_args.experiment == 'random':
     #     torch.save(model.state_dict(), path_save)
 
-
-    if data_args.experiment_mode == 'lm' and data_args.modality in ['roc-aug', 'roc', 'yelp', 'commonGen', 'commonGen-aug'] \
-            and data_args.cache_mode=='no':
-        train_dataset = helper_tokenize_stream(sentence_lst, vocab_dict, model, image_size**2, data_args, padding_mode)
+    if (
+        data_args.experiment_mode == "lm"
+        and data_args.modality
+        in ["roc-aug", "roc", "yelp", "commonGen", "commonGen-aug"]
+        and data_args.cache_mode == "no"
+    ):
+        train_dataset = helper_tokenize_stream(
+            sentence_lst, vocab_dict, model, image_size**2, data_args, padding_mode
+        )
         return train_dataset, model
-    elif data_args.experiment_mode == 'lm':
-        result_train_lst = helper_tokenize_encode(sentence_lst, vocab_dict, model, image_size**2, data_args, padding_mode)
-    elif data_args.experiment_mode == 'conditional_gen':
-        result_train_lst = helper_tokenize_encode_cond(sentence_lst, vocab_dict, model, image_size ** 2, data_args)
-    return {'train': result_train_lst}, model
-       
+    elif data_args.experiment_mode == "lm":
+        result_train_lst = helper_tokenize_encode(
+            sentence_lst, vocab_dict, model, image_size**2, data_args, padding_mode
+        )
+    elif data_args.experiment_mode == "conditional_gen":
+        result_train_lst = helper_tokenize_encode_cond(
+            sentence_lst, vocab_dict, model, image_size**2, data_args
+        )
+    return {"train": result_train_lst}, model
+
 
 def write_e2e_corr(prompt_lst, file_dict, corr_path):
     print(len(prompt_lst))
-    with open(corr_path, 'w') as f:
+    with open(corr_path, "w") as f:
         for x in prompt_lst:
             for line in file_dict[x]:
                 print(" ".join(line), file=f)
-            print('', file=f)
+            print("", file=f)
 
 
 def write_e2e_src(prompt_lst, corr_path):
-    with open(corr_path, 'w') as f:
+    with open(corr_path, "w") as f:
         for x in prompt_lst:
             print(" ".join(x), file=f)
     return
@@ -624,48 +710,55 @@ def write_e2e_src(prompt_lst, corr_path):
 
 def read_e2e_files(path, args, tokenizer):
     file_dict = {}
-    with open(path, 'r') as f:
+    with open(path, "r") as f:
         for line in f:
-            src_lst, word_lst = line.strip().split('||')
+            src_lst, word_lst = line.strip().split("||")
             tgt = tuple([x.text for x in tokenizer(word_lst)])
             src = tuple([x.text for x in tokenizer(src_lst)])
             if src not in file_dict:
                 file_dict[src] = []
             file_dict[src].append(tgt)
-    temp = '1'
+    temp = "1"
     prompt_text_dict = file_dict
     prompt_text_lst = list(prompt_text_dict.keys())
-    gold_dir = os.path.join(args.out_dir, '{}_{}_{}'.format(temp, args.split, 'gold'))
+    gold_dir = os.path.join(args.out_dir, "{}_{}_{}".format(temp, args.split, "gold"))
     print("gold dir", gold_dir)
     write_e2e_corr(prompt_text_lst, prompt_text_dict, gold_dir)
-    src_dir = os.path.join(args.out_dir, '{}_{}_{}'.format(temp, args.split, 'src'))
+    src_dir = os.path.join(args.out_dir, "{}_{}_{}".format(temp, args.split, "src"))
     write_e2e_src(prompt_text_lst, src_dir)
     final_lst = [(xx, prompt_text_dict[xx][0]) for xx in prompt_text_lst]
     return final_lst
 
 
-def get_corpus_book(data_args, tokenizer, model, image_size, padding_mode='block', split='train',):
-    max_length = image_size ** 2
+def get_corpus_book(
+    data_args,
+    tokenizer,
+    model,
+    image_size,
+    padding_mode="block",
+    split="train",
+):
+    max_length = image_size**2
     import os
-    assert padding_mode == 'block'
-    raw_datasets = load_dataset('bookcorpus')
+
+    assert padding_mode == "block"
+    raw_datasets = load_dataset("bookcorpus")
     if "validation" not in raw_datasets.keys():
         raw_datasets["validation"] = load_dataset(
-            'bookcorpus',
+            "bookcorpus",
             split=f"train[:1%]",
         )
         raw_datasets["train"] = load_dataset(
-            'bookcorpus',
+            "bookcorpus",
             split=f"train[1%:]",
         )
     print(raw_datasets)
     column_names = raw_datasets["train"].column_names
 
     def tokenize_function(examples):
-        output = tokenizer(examples['text'], add_special_tokens=False)
+        output = tokenizer(examples["text"], add_special_tokens=False)
         return output
 
-
     tokenized_datasets = raw_datasets.map(
         tokenize_function,
         batched=True,
@@ -686,7 +779,7 @@ def get_corpus_book(data_args, tokenizer, model, image_size, padding_mode='block
         if total_length >= block_size:
             total_length = (total_length // block_size) * block_size
         result = {
-            k: [t[i: i + block_size] for i in range(0, total_length, block_size)]
+            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
             for k, t in concatenated_examples.items()
         }
         return result
@@ -702,32 +795,44 @@ def get_corpus_book(data_args, tokenizer, model, image_size, padding_mode='block
     print(lm_datasets)
 
     if model is None:
-        if data_args.training_mode.startswith('e2e'):
-            print('since its e2e, initialize a dummy embedding' )
+        if data_args.training_mode.startswith("e2e"):
+            print("since its e2e, initialize a dummy embedding")
             model = torch.nn.Embedding(len(tokenizer), 1)
         else:
             model = torch.nn.Embedding(len(tokenizer), data_args.in_channel)
-        print('initializing the random embeddings', model)
+        print("initializing the random embeddings", model)
         torch.nn.init.normal_(model.weight)
-        path_save = f'{data_args.checkpoint_path}/random_emb.torch'
-        print(f'save the random encoder to {data_args.checkpoint_path}/random_emb.torch')
+        path_save = f"{data_args.checkpoint_path}/random_emb.torch"
+        print(
+            f"save the random encoder to {data_args.checkpoint_path}/random_emb.torch"
+        )
         torch.save(model.state_dict(), path_save)
 
-    if split == 'train':
+    if split == "train":
         return lm_datasets, model
     else:
-        lm_datasets['train'] = lm_datasets['validation']
+        lm_datasets["train"] = lm_datasets["validation"]
         return lm_datasets, model
 
 
 class TextDataset(Dataset):
-    def __init__(self, text_datasets, resolution, data_args, model_arch='conv-unet',
-                 classes=None, shard=0, num_shards=1, eigen_transform=None,
-                 mapping_func=None, model_emb=None):
+    def __init__(
+        self,
+        text_datasets,
+        resolution,
+        data_args,
+        model_arch="conv-unet",
+        classes=None,
+        shard=0,
+        num_shards=1,
+        eigen_transform=None,
+        mapping_func=None,
+        model_emb=None,
+    ):
         super().__init__()
         self.resolution = resolution
         self.text_datasets = text_datasets
-        self.length = len(self.text_datasets['train'])
+        self.length = len(self.text_datasets["train"])
         self.model_arch = model_arch
         self.data_args = data_args
         print(self.resolution)
@@ -745,8 +850,8 @@ class TextDataset(Dataset):
         # We are not on a new enough PIL to support the `reducing_gap`
         # argument, which uses BOX downsampling at powers of two first.
         # Thus, we do it by hand to improve downsample quality.
-        if self.model_arch == 'conv-unet':
-            pass# arr = np.array(self.text_datasets['train'][idx]['hidden_states'],
+        if self.model_arch == "conv-unet":
+            pass  # arr = np.array(self.text_datasets['train'][idx]['hidden_states'],
             #                dtype=np.float32).reshape(self.resolution, self.resolution, -1)
             # # print(self.eigen_transform.shape)
             # if self.eigen_transform  is not None:
@@ -757,15 +862,14 @@ class TextDataset(Dataset):
             # if hasattr(self.data_args, 'noise_level') and self.data_args.noise_level > 0:
             #     arr = arr + self.data_args.noise_level * np.random.randn(*arr.shape).astype(arr.dtype)
 
-
             # out_dict = {}
             # out_dict['input_ids'] = np.array(self.text_datasets['train'][idx]['input_ids'])
             # # if self.local_classes is not None:
             # #     out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64)
             # # print(out_dict.keys())
             # return np.transpose(arr, [2, 0, 1]), out_dict
-        elif self.model_arch == '1d-unet':
-            pass# arr = np.array(self.text_datasets['train'][idx]['hidden_states'],
+        elif self.model_arch == "1d-unet":
+            pass  # arr = np.array(self.text_datasets['train'][idx]['hidden_states'],
             #                dtype=np.float32) # seqlen, dim
             # if self.eigen_transform  is not None:
             #     old_shape = arr.shape
@@ -783,27 +887,39 @@ class TextDataset(Dataset):
             # # print(arr.shape)
             # return arr, out_dict
         else:
-            arr = np.array(self.text_datasets['train'][idx]['hidden_states'],
-                           dtype=np.float32)
-            if self.eigen_transform  is not None:
+            arr = np.array(
+                self.text_datasets["train"][idx]["hidden_states"], dtype=np.float32
+            )
+            if self.eigen_transform is not None:
                 old_shape = arr.shape
                 # arr = arr.reshape(1, -1) @ self.eigen_transform
-                arr = arr.reshape(1, -1) - self.eigen_transform['mean']
-                arr = arr @ self.eigen_transform['map']
+                arr = arr.reshape(1, -1) - self.eigen_transform["mean"]
+                arr = arr @ self.eigen_transform["map"]
                 arr = arr.reshape(old_shape)
-                
-            if hasattr(self.data_args, 'noise_level') and self.data_args.noise_level > 0:
+
+            if (
+                hasattr(self.data_args, "noise_level")
+                and self.data_args.noise_level > 0
+            ):
                 # print(arr.dtype)
                 # print(self.data_args.noise_level, 'using the noise level.')
-                arr = arr + self.data_args.noise_level * np.random.randn(*arr.shape).astype(arr.dtype)
+                arr = arr + self.data_args.noise_level * np.random.randn(
+                    *arr.shape
+                ).astype(arr.dtype)
                 # print(arr.dtype)
 
             out_dict = {}
-            out_dict['input_ids'] = np.array(self.text_datasets['train'][idx]['input_ids'])
+            out_dict["input_ids"] = np.array(
+                self.text_datasets["train"][idx]["input_ids"]
+            )
             # out_dict['mapping_func'] = self.mapping_func
-            if self.data_args.experiment_mode == 'conditional_gen':
-                out_dict['src_ids'] = np.array(self.text_datasets['train'][idx]['src_ids'])
-                out_dict['src_mask'] = np.array(self.text_datasets['train'][idx]['src_mask'])
+            if self.data_args.experiment_mode == "conditional_gen":
+                out_dict["src_ids"] = np.array(
+                    self.text_datasets["train"][idx]["src_ids"]
+                )
+                out_dict["src_mask"] = np.array(
+                    self.text_datasets["train"][idx]["src_mask"]
+                )
             # if self.local_classes is not None:
             #     out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64)
             return arr, out_dict
@@ -813,13 +929,23 @@ class TextDataset(Dataset):
 
 
 class TextDataset_NoCache(Dataset):
-    def __init__(self, text_datasets, resolution, data_args, model_arch='conv-unet',
-                 classes=None, shard=0, num_shards=1, eigen_transform=None,
-                 mapping_func=None, model_emb=None):
+    def __init__(
+        self,
+        text_datasets,
+        resolution,
+        data_args,
+        model_arch="conv-unet",
+        classes=None,
+        shard=0,
+        num_shards=1,
+        eigen_transform=None,
+        mapping_func=None,
+        model_emb=None,
+    ):
         super().__init__()
         self.resolution = resolution
         self.text_datasets = text_datasets
-        self.length = len(self.text_datasets['train'])
+        self.length = len(self.text_datasets["train"])
         self.model_arch = model_arch
         self.data_args = data_args
         print(self.resolution)
@@ -838,81 +964,110 @@ class TextDataset_NoCache(Dataset):
         # argument, which uses BOX downsampling at powers of two first.
         # Thus, we do it by hand to improve downsample quality.
         with torch.no_grad():
-            input_ids = self.text_datasets['train'][idx]['input_ids']
+            input_ids = self.text_datasets["train"][idx]["input_ids"]
             model = self.model_emb
-            if self.data_args.experiment.startswith('random'):
+            if self.data_args.experiment.startswith("random"):
                 hidden_state = model(torch.tensor(input_ids))
-            elif self.data_args.experiment == 'gpt2_pre_compress':
+            elif self.data_args.experiment == "gpt2_pre_compress":
                 input_ids2 = torch.tensor(input_ids).to(model.device)
                 input_embs = model.transformer.wte(input_ids2)  # input_embs
                 hidden_state = model.down_proj(input_embs)
                 hidden_state = hidden_state * data_args.emb_scale_factor
 
-            if self.model_arch == 'conv-unet':
-                arr = np.array(hidden_state,
-                               dtype=np.float32).reshape(self.resolution, self.resolution, -1)
+            if self.model_arch == "conv-unet":
+                arr = np.array(hidden_state, dtype=np.float32).reshape(
+                    self.resolution, self.resolution, -1
+                )
                 # print(self.eigen_transform.shape)
                 if self.eigen_transform is not None:
                     old_shape = arr.shape
-                    arr = arr.reshape(1, -1) - self.eigen_transform['mean']
-                    arr = arr @ self.eigen_transform['map']
+                    arr = arr.reshape(1, -1) - self.eigen_transform["mean"]
+                    arr = arr @ self.eigen_transform["map"]
                     arr = arr.reshape(old_shape)
-                if hasattr(self.data_args, 'noise_level') and self.data_args.noise_level > 0:
-                    arr = arr + self.data_args.noise_level * np.random.randn(*arr.shape).astype(arr.dtype)
+                if (
+                    hasattr(self.data_args, "noise_level")
+                    and self.data_args.noise_level > 0
+                ):
+                    arr = arr + self.data_args.noise_level * np.random.randn(
+                        *arr.shape
+                    ).astype(arr.dtype)
 
                 out_dict = {}
-                out_dict['input_ids'] = np.array(self.text_datasets['train'][idx]['input_ids'])
+                out_dict["input_ids"] = np.array(
+                    self.text_datasets["train"][idx]["input_ids"]
+                )
                 # if self.local_classes is not None:
                 #     out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64)
                 # print(out_dict.keys())
                 return np.transpose(arr, [2, 0, 1]), out_dict
-            elif self.model_arch == '1d-unet':
-                arr = np.array(hidden_state,
-                               dtype=np.float32)  # seqlen, dim
+            elif self.model_arch == "1d-unet":
+                arr = np.array(hidden_state, dtype=np.float32)  # seqlen, dim
                 if self.eigen_transform is not None:
                     old_shape = arr.shape
-                    arr = arr.reshape(1, -1) - self.eigen_transform['mean']
-                    arr = arr @ self.eigen_transform['map']
+                    arr = arr.reshape(1, -1) - self.eigen_transform["mean"]
+                    arr = arr @ self.eigen_transform["map"]
                     arr = arr.reshape(old_shape)
-                if hasattr(self.data_args, 'noise_level') and self.data_args.noise_level > 0:
-                    arr = arr + self.data_args.noise_level * np.random.randn(*arr.shape).astype(arr.dtype)
+                if (
+                    hasattr(self.data_args, "noise_level")
+                    and self.data_args.noise_level > 0
+                ):
+                    arr = arr + self.data_args.noise_level * np.random.randn(
+                        *arr.shape
+                    ).astype(arr.dtype)
                 arr = np.transpose(arr, [1, 0])
                 out_dict = {}
-                out_dict['input_ids'] = np.array(self.text_datasets['train'][idx]['input_ids'])
+                out_dict["input_ids"] = np.array(
+                    self.text_datasets["train"][idx]["input_ids"]
+                )
                 # out_dict['mapping_func'] = self.mapping_func
                 # if self.local_classes is not None:
                 #     out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64)
                 # print(arr.shape)
                 return arr, out_dict
             else:
-                arr = np.array(hidden_state,
-                               dtype=np.float32)
+                arr = np.array(hidden_state, dtype=np.float32)
                 if self.eigen_transform is not None:
                     old_shape = arr.shape
                     # arr = arr.reshape(1, -1) @ self.eigen_transform
-                    arr = arr.reshape(1, -1) - self.eigen_transform['mean']
-                    arr = arr @ self.eigen_transform['map']
+                    arr = arr.reshape(1, -1) - self.eigen_transform["mean"]
+                    arr = arr @ self.eigen_transform["map"]
                     arr = arr.reshape(old_shape)
 
-                if hasattr(self.data_args, 'noise_level') and self.data_args.noise_level > 0:
+                if (
+                    hasattr(self.data_args, "noise_level")
+                    and self.data_args.noise_level > 0
+                ):
                     # print(arr.dtype)
                     # print(self.data_args.noise_level, 'using the noise level.')
-                    arr = arr + self.data_args.noise_level * np.random.randn(*arr.shape).astype(arr.dtype)
+                    arr = arr + self.data_args.noise_level * np.random.randn(
+                        *arr.shape
+                    ).astype(arr.dtype)
                     # print(arr.dtype)
 
                 out_dict = {}
-                out_dict['input_ids'] = np.array(self.text_datasets['train'][idx]['input_ids'])
+                out_dict["input_ids"] = np.array(
+                    self.text_datasets["train"][idx]["input_ids"]
+                )
                 # out_dict['mapping_func'] = self.mapping_func
-                if self.data_args.experiment_mode == 'conditional_gen':
-                    out_dict['src_ids'] = np.array(self.text_datasets['train'][idx]['src_ids'])
-                    out_dict['src_mask'] = np.array(self.text_datasets['train'][idx]['src_mask'])
+                if self.data_args.experiment_mode == "conditional_gen":
+                    out_dict["src_ids"] = np.array(
+                        self.text_datasets["train"][idx]["src_ids"]
+                    )
+                    out_dict["src_mask"] = np.array(
+                        self.text_datasets["train"][idx]["src_mask"]
+                    )
                 # if self.local_classes is not None:
                 #     out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64)
                 return arr, out_dict
 
+
 def _collate_batch_helper(examples, pad_token_id, max_length, return_mask=False):
-    result = torch.full([len(examples), max_length], pad_token_id, dtype=torch.int64).tolist()
-    mask_ = torch.full([len(examples), max_length], pad_token_id, dtype=torch.int64).tolist()
+    result = torch.full(
+        [len(examples), max_length], pad_token_id, dtype=torch.int64
+    ).tolist()
+    mask_ = torch.full(
+        [len(examples), max_length], pad_token_id, dtype=torch.int64
+    ).tolist()
     for i, example in enumerate(examples):
         curr_len = min(len(example), max_length)
         result[i][:curr_len] = example[:curr_len]
@@ -921,6 +1076,7 @@ def _collate_batch_helper(examples, pad_token_id, max_length, return_mask=False)
         return result, mask_
     return result
 
+
 def _torch_collate_batch(examples, pad_token_id, max_length):
     """Collate `examples` into a batch, using the information in `tokenizer` for padding if necessary."""
     import numpy as np
@@ -945,4 +1101,4 @@ def _torch_collate_batch(examples, pad_token_id, max_length):
             result[i, : example.shape[0]] = example
         else:
             result[i, -example.shape[0] :] = example
-    return result
+    return result