src/improved_diffusion/text_datasets.py

# from PIL import Image
# import blobfile as bf
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 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
from functools import partial
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,
):
    """
    For a dataset, create a generator over (images, kwargs) pairs.

    Each images is an NCHW float tensor, and the kwargs dict contains zero or
    more keys, each of which map to a batched Tensor of their own.
    The kwargs dict can be used for class labels, in which case the key is "y"
    and the values are integer tensors of class labels.

    :param data_dir: a dataset directory.
    :param batch_size: the batch size of each returned pair.
    :param image_size: the size to which images are resized.
    :param class_cond: if True, include a "y" key in returned dicts for class
                       label. If classes are not available and this is true, an
                       exception will be raised.
    :param deterministic: if True, yield results in a deterministic order.
    """
    print('hello loading text data. ')

    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' :
        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':
        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)
    # elif task_mode == 'yelp':
    #     print('hello loading yelp ')
    #     training_data, model = get_corpus_rocstory(data_args, model, image_size,
    #                                         padding_mode=padding_mode, split=split,
    #                                         load_vocab=load_vocab)

    # elif task_mode == 'commonGen' or task_mode == 'commonGen-aug':
    #     print('hello loading common-gen ')
    #     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':
    #     training_data, model = get_corpus_rocstory(data_args, model, image_size,
    #                                         padding_mode=padding_mode, split=split,
    #                                         load_vocab=load_vocab)

    # elif task_mode == 'book':
    #     tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
    #     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(
        #     training_data,
        #     image_size,
        #     data_args,
        #     model_arch=data_args.model_arch,
        #     model_emb=model
        # )
    else:
        dataset = TextDataset(
            training_data,
            image_size,
            data_args,
            model_arch=data_args.model_arch,
        )

    if deterministic:

        pass# data_loader = DataLoader(
        #     dataset,
        #     batch_size=batch_size,  # 20,
        #     drop_last=True,
        #     shuffle=False,
        #     num_workers=1,
        # )

    else:
        data_loader = DataLoader(
            dataset,
            batch_size=batch_size,  # 20,
            drop_last=True,
            shuffle=True,
            num_workers=1,
        )
    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]
            input_ids = [0] + tokenized_ + [1]
            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')
        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']])
        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)


        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':
                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
                                     })

    return result_train_lst

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})
    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']]
        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}
        # clm input could be much much longer than block_size
        return result_dict

    tokenized_datasets = raw_datasets.map(
        tokenize_function,
        batched=True,
        num_proc=4,
        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':
        block_size = seqlen
        def group_texts(examples):
            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)]
                for k, t in concatenated_examples.items()
            }
            result["labels"] = result["input_ids"].copy()
            return result


        lm_datasets = tokenized_datasets.map(
            group_texts,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            load_from_cache_file=not data_args.overwrite_cache,
            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)
            else:
                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
        print(f"RAM used: {psutil.Process().memory_info().rss / (1024 * 1024):.2f} MB")

        lm_datasets = tokenized_datasets.map(
            pad_function,
            batched=True,
            num_proc=1,
            desc=f"padding",
        )


    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
    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, ):
    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]
            input_ids = [0] + tokenized_ + [1]
            group_lst['word_ids'].append(input_ids)
        print(group_lst['word_ids'][:2])

        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)]
                for k, t in concatenated_examples.items()
            }
        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)

        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':
                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':
                hidden_state = model(torch.tensor(input_ids))
            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:
        for line in f:
            split_line = line.split()
            word = split_line[0]
            embedding = torch.tensor(np.array(split_line[1:], dtype=np.float64))
            # embedding = np.array(split_line[1:], dtype=np.float64)
            glove_model[word] = embedding
    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')
    array_lst = []
    count_ = 0
    for word, idx in vocab.items():
        if word in glove_model:
            array_lst.append(glove_model[word])
        else:
            count_ += 1
            array_lst.append(torch.randn(50))
    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):
    import csv, torch, json
    from spacy.lang.en import English

    if data_args.experiment_mode == 'lm':
        if data_args.modality == 'roc':
            pass
            # print('loading dataset from ROCStory')
            # nlp = English()
            # tokenizer = nlp.tokenizer
            # sentence_lst = []
            # print(f'loading from {data_args.roc_train}')
            # if split == 'train':
            #     print('loading form the TRAIN set')
            #     path = f'{data_args.roc_train}/roc_train.json'
            # elif split == 'valid':
            #     print('loading form the VALID set')
            #     path = f'{data_args.roc_train}/roc_valid.json'
            # else:
            #     assert False, "invalid split for ROC dataset"

            # 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)]
            #         sentence_lst.append(word_lst)

            # # with open(data_args.roc_train, 'r') as csvfile:
            # #     roc_reader = csv.reader(csvfile) #delimiter=' ', quotechar='|')
            # #     for row in roc_reader:
            # #         # tokenize.
            # #         sentences = " ".join(row[2:])
            # #         word_lst = [x.text for x in tokenizer(sentences)]
            # #         sentence_lst.append(word_lst)
            # # sentence_lst = sentence_lst[1:]
            # print(sentence_lst[:2])
        if data_args.modality == 'roc-aug':
            pass
            # print('loading dataset from ROCStory')
            # nlp = English()
            # tokenizer = nlp.tokenizer
            # sentence_lst = []
            # if split == 'train':
            #     print('loading form the TRAIN set')
            #     path_lst = [f'{data_args.roc_train}/roc_train.json']
            #     path_lst.append('diffusion_lm/improved-diffusion/diff_models/rocstories_gptj.txt')
            #     # path_lst.append('diffusion_lm/improved-diffusion/cache/ar_model_augment_roc.json')
            #     # path_lst.append('diffusion_lm/improved-diffusion/cache/ar_model_augment_roc2.json')

            # elif split == 'valid':
            #     print('loading form the VALID set')
            #     path_lst = [f'{data_args.roc_train}/roc_valid.json']
            # else:
            #     assert False, "invalid split for ROC dataset"

            # print(path_lst)
            # for path in path_lst:
            #     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)]
            #                 sentence_lst.append(word_lst)
            #     else:
            #         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)]
            #                 sentence_lst.append(word_lst)
            # print(sentence_lst[:2],sentence_lst[-2:], 'dataset size=',len(sentence_lst))
        elif data_args.modality == 'simple-wiki':
            pass
            # print('loading dataset from simple wikipedia')
            # sentence_lst = []
            # with open(data_args.wiki_train, 'r') as ff:
            #     for row in ff:
            #         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')
            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'
                # 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')
                path = data_args.debug_path
                import json
                with open(path, 'r') as ff:
                    for line in ff:
                        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:
                    for row in ff:
                        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')
            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='|')
                    for row in yelp_reader:
                        sentences = row[1]
                        word_lst = [x.text for x in tokenizer(sentences)]
                        sentence_lst.append(word_lst)
                sentence_lst = sentence_lst[1:]
            print(sentence_lst[:2])

        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:
                    for line in ff:
                        line = json.loads(line)
                        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')
            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'
                path_lst = []
            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:
                    for line in ff:
                        line = json.loads(line)
                        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:
                        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 == '.'):
                                spl[-1].extend(y)
                                if x: spl.append([])
                            sentence_lst.extend(spl[:-1])
                else:
                    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 == '.'):
                                spl[-1].extend(y)
                                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')
            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:
                    for row in ff:
                        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'
                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:
                counter.update(input_ids)
                counter.update(src_ids)

    if load_vocab is None:
        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:
            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'
        if not os.path.exists(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:
                    json.dump(vocab_dict, f)
                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':
        model = torch.nn.Embedding(len(vocab_dict), data_args.in_channel)
        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')
        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)
        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
       

def write_e2e_corr(prompt_lst, file_dict, corr_path):
    print(len(prompt_lst))
    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)


def write_e2e_src(prompt_lst, corr_path):
    with open(corr_path, 'w') as f:
        for x in prompt_lst:
            print(" ".join(x), file=f)
    return


def read_e2e_files(path, args, tokenizer):
    file_dict = {}
    with open(path, 'r') as f:
        for line in f:
            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'
    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'))
    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'))
    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
    import os
    assert padding_mode == 'block'
    raw_datasets = load_dataset('bookcorpus')
    if "validation" not in raw_datasets.keys():
        raw_datasets["validation"] = load_dataset(
            'bookcorpus',
            split=f"train[:1%]",
        )
        raw_datasets["train"] = load_dataset(
            '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)
        return output


    tokenized_datasets = raw_datasets.map(
        tokenize_function,
        batched=True,
        num_proc=data_args.preprocessing_num_workers,
        remove_columns=column_names,
        load_from_cache_file=True,
    )

    print(tokenized_datasets)

    block_size = max_length

    # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
    def group_texts(examples):
        # Concatenate all texts.
        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)]
            for k, t in concatenated_examples.items()
        }
        return result

    lm_datasets = tokenized_datasets.map(
        group_texts,
        batched=True,
        num_proc=4,
        load_from_cache_file=True,
        desc=f"Grouping texts in chunks of {block_size}",
    )

    print(lm_datasets)

    if model is None:
        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)
        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')
        torch.save(model.state_dict(), path_save)

    if split == 'train':
        return lm_datasets, model
    else:
        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):
        super().__init__()
        self.resolution = resolution
        self.text_datasets = text_datasets
        self.length = len(self.text_datasets['train'])
        self.model_arch = model_arch
        self.data_args = data_args
        print(self.resolution)
        self.eigen_transform = eigen_transform
        self.mapping_func = mapping_func
        self.model_emb = model_emb
        # self.local_images = image_paths[shard:][::num_shards]
        # self.local_classes = None if classes is None else classes[shard:][::num_shards]

    def __len__(self):
        return self.length

    def __getitem__(self, idx):

        # 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'],
            #                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(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)


            # 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'],
            #                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(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)
            # arr = np.transpose(arr, [1, 0])
            # out_dict = {}
            # 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(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(old_shape)
                
            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)
                # print(arr.dtype)

            out_dict = {}
            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.local_classes is not None:
            #     out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64)
            return arr, out_dict
        # print(arr.dtype)
        # arr = arr.float()
        # print(arr.shape)


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):
        super().__init__()
        self.resolution = resolution
        self.text_datasets = text_datasets
        self.length = len(self.text_datasets['train'])
        self.model_arch = model_arch
        self.data_args = data_args
        print(self.resolution)
        self.eigen_transform = eigen_transform
        self.mapping_func = mapping_func
        self.model_emb = model_emb
        # self.local_images = image_paths[shard:][::num_shards]
        # self.local_classes = None if classes is None else classes[shard:][::num_shards]

    def __len__(self):
        return self.length

    def __getitem__(self, idx):

        # 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.
        with torch.no_grad():
            input_ids = self.text_datasets['train'][idx]['input_ids']
            model = self.model_emb
            if self.data_args.experiment.startswith('random'):
                hidden_state = model(torch.tensor(input_ids))
            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)
                # 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(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)

                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':
                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(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)
                arr = np.transpose(arr, [1, 0])
                out_dict = {}
                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)
                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(old_shape)

                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)
                    # print(arr.dtype)

                out_dict = {}
                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.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()
    for i, example in enumerate(examples):
        curr_len = min(len(example), max_length)
        result[i][:curr_len] = example[:curr_len]
        mask_[i][:curr_len] = [1] * curr_len
    if return_mask:
        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
    import torch

    # Tensorize if necessary.
    if isinstance(examples[0], (list, tuple, np.ndarray)):
        examples = [torch.tensor(e, dtype=torch.long) for e in examples]

    # length_of_first = examples[0].size(0)
    # Check if padding is necessary.
    # are_tensors_same_length = all(x.size(0) == length_of_first for x in examples)
    # if are_tensors_same_length and (pad_to_multiple_of is None or length_of_first % pad_to_multiple_of == 0):
    #     return torch.stack(examples, dim=0)
    # Creating the full tensor and filling it with our data.
    # max_length = max(x.size(0) for x in examples)
    # if pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
    #     max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of
    result = examples[0].new_full([len(examples), max_length], pad_token_id)
    for i, example in enumerate(examples):
        if True:
            result[i, : example.shape[0]] = example
        else:
            result[i, -example.shape[0] :] = example
    return result