hashed_index.py 5.53 KB
Newer Older
Neel Kant's avatar
Neel Kant committed
1
from collections import defaultdict
Neel Kant's avatar
Neel Kant committed
2
import os
Neel Kant's avatar
Neel Kant committed
3
import pickle
Neel Kant's avatar
Neel Kant committed
4
import shutil
Neel Kant's avatar
Neel Kant committed
5

Neel Kant's avatar
Neel Kant committed
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import numpy as np
import torch
from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP

from megatron import get_args
from megatron import mpu
from megatron.checkpointing import get_checkpoint_tracker_filename, get_checkpoint_name
from megatron.data.bert_dataset import get_indexed_dataset_
from megatron.data.ict_dataset import InverseClozeDataset
from megatron.data.samplers import DistributedBatchSampler
from megatron.initialize import initialize_megatron
from megatron.training import get_model
from pretrain_bert_ict import get_batch, model_provider


21
22
23
24
def detach(tensor):
    return tensor.detach().cpu().numpy()


Neel Kant's avatar
Neel Kant committed
25
def embed_docs():
Neel Kant's avatar
Neel Kant committed
26
27
28
29
30
31
32
33
    initialize_megatron(extra_args_provider=None,
                        args_defaults={'tokenizer_type': 'BertWordPieceLowerCase'})
    args = get_args()
    model = load_checkpoint()
    model.eval()
    dataset = get_dataset()
    data_iter = iter(get_dataloader(dataset))

Neel Kant's avatar
Neel Kant committed
34
    hash_data = defaultdict(list)
Neel Kant's avatar
Neel Kant committed
35
    hash_matrix = torch.cuda.HalfTensor(np.random.rand(128, 1024))
Neel Kant's avatar
Neel Kant committed
36
    hash_data['matrix'] = hash_matrix
Neel Kant's avatar
Neel Kant committed
37

38
    block_data = defaultdict(list)
Neel Kant's avatar
Neel Kant committed
39
    i = 0
Neel Kant's avatar
Neel Kant committed
40
41
42
43
    while True:
        try:
            input_tokens, input_types, input_pad_mask, \
            block_tokens, block_token_types, block_pad_mask, block_indices = get_batch(data_iter)
44
        except:
Neel Kant's avatar
Neel Kant committed
45
            break
46

Neel Kant's avatar
Neel Kant committed
47
48
        input_logits, block_logits = model.module.module.forward(
            input_tokens, input_types, input_pad_mask, block_tokens, block_pad_mask, block_token_types)
Neel Kant's avatar
Neel Kant committed
49

Neel Kant's avatar
Neel Kant committed
50
        block_hash_pos = torch.matmul(block_logits, hash_matrix)
Neel Kant's avatar
Neel Kant committed
51
        block_hash_full = torch.cat((block_hash_pos, -block_hash_pos), axis=1)
52
        block_hashes = detach(torch.argmax(block_hash_full, axis=1))
53
        for hash, indices_array in zip(block_hashes, block_indices):
54
            hash_data[int(hash)].append(detach(indices_array))
Neel Kant's avatar
Neel Kant committed
55

56
57
58
        block_logits = detach(block_logits)
        # originally this has [start_idx, end_idx, doc_idx, block_idx]
        block_indices = detach(block_indices)[:, 3]
59
        for logits, idx in zip(block_logits, block_indices):
Neel Kant's avatar
Neel Kant committed
60
            block_data[int(idx)] = logits
Neel Kant's avatar
Neel Kant committed
61

Neel Kant's avatar
Neel Kant committed
62
63
        if i % 100 == 0:
            print(i, flush=True)
64
65
        i += 1

Neel Kant's avatar
Neel Kant committed
66
67
68
    dir_name = 'block_hash_data'
    if not os.path.isdir(dir_name):
        os.mkdir(dir_name)
69

Neel Kant's avatar
Neel Kant committed
70
    # save the data for each shard
Neel Kant's avatar
Neel Kant committed
71
72
73
    with open('{}/{}.pkl'.format(dir_name, args.rank), 'wb') as data_file:
        all_data = {'block_data': block_data, 'hash_data': hash_data}
        pickle.dump(all_data, data_file)
74

Neel Kant's avatar
Neel Kant committed
75
    torch.distributed.barrier()
Neel Kant's avatar
Neel Kant committed
76

77
78
79
80
    all_data.clear()
    del all_data
    del model

Neel Kant's avatar
Neel Kant committed
81
    # rank 0 process consolidates shards and saves into final file
Neel Kant's avatar
Neel Kant committed
82
83
84
85
86
    if mpu.get_data_parallel_rank() == 0:
        all_block_data = defaultdict(dict)
        dir_name = 'block_hash_data'
        fnames = os.listdir(dir_name)
        for fname in fnames:
Neel Kant's avatar
Neel Kant committed
87
            with open('{}/{}'.format(dir_name, fname), 'rb') as f:
Neel Kant's avatar
Neel Kant committed
88
89
90
                data = pickle.load(f)
                all_block_data['hash_data'].update(data['hash_data'])
                all_block_data['block_data'].update(data['block_data'])
Neel Kant's avatar
Neel Kant committed
91

Neel Kant's avatar
Neel Kant committed
92
93
        with open('block_hash_data.pkl', 'wb') as final_file:
            pickle.dump(all_block_data, final_file)
Neel Kant's avatar
Neel Kant committed
94
        shutil.rmtree(dir_name, ignore_errors=True)
Neel Kant's avatar
Neel Kant committed
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124


def load_checkpoint():
    args = get_args()
    model = get_model(model_provider)

    if isinstance(model, torchDDP):
        model = model.module
    tracker_filename = get_checkpoint_tracker_filename(args.load)
    with open(tracker_filename, 'r') as f:
        iteration = int(f.read().strip())

    assert iteration > 0
    checkpoint_name = get_checkpoint_name(args.load, iteration, False)
    if mpu.get_data_parallel_rank() == 0:
        print('global rank {} is loading checkpoint {}'.format(
            torch.distributed.get_rank(), checkpoint_name))

    state_dict = torch.load(checkpoint_name, map_location='cpu')
    model.load_state_dict(state_dict['model'])
    torch.distributed.barrier()

    if mpu.get_data_parallel_rank() == 0:
        print(' successfully loaded {}'.format(checkpoint_name))

    return model


def get_dataset():
    args = get_args()
Neel Kant's avatar
Neel Kant committed
125
126
    block_dataset = get_indexed_dataset_(args.data_path, 'mmap', True)
    titles_dataset = get_indexed_dataset_(args.data_path + '-titles', 'mmap', True)
Neel Kant's avatar
Neel Kant committed
127
128
129

    kwargs = dict(
        name='full',
Neel Kant's avatar
Neel Kant committed
130
131
        block_dataset=block_dataset,
        title_dataset=titles_dataset,
Neel Kant's avatar
Neel Kant committed
132
        data_prefix=args.data_path,
Neel Kant's avatar
Neel Kant committed
133
134
        num_epochs=1,
        max_num_samples=None,
Neel Kant's avatar
Neel Kant committed
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
        max_seq_length=288,  # doesn't matter
        short_seq_prob=0.0001,  # doesn't matter
        seed=1
    )
    dataset = InverseClozeDataset(**kwargs)
    return dataset


def get_dataloader(dataset):
    args = get_args()

    world_size = mpu.get_data_parallel_world_size()
    rank = mpu.get_data_parallel_rank()
    global_batch_size = args.batch_size * world_size
    num_workers = args.num_workers

    sampler = torch.utils.data.SequentialSampler(dataset)
    batch_sampler = DistributedBatchSampler(sampler,
                                            batch_size=global_batch_size,
                                            drop_last=True,
                                            rank=rank,
                                            world_size=world_size)

    return torch.utils.data.DataLoader(dataset,
                                       batch_sampler=batch_sampler,
                                       num_workers=num_workers,
                                       pin_memory=True)


if __name__ == "__main__":
Neel Kant's avatar
Neel Kant committed
165
    embed_docs()