extractor.py

# coding=utf-8
# Copyright (c) Microsoft. All rights reserved.
"""Extract feature vectors.
"""
import os
import argparse
import torch
import json
from pytorch_pretrained_bert.tokenization import BertTokenizer
from torch.utils.data import DataLoader
from data_utils.log_wrapper import create_logger
from data_utils.utils import set_environment
from mt_dnn.batcher import Collater, SingleTaskDataset
from mt_dnn.model import MTDNNModel
from prepro_std import _truncate_seq_pair
from data_utils.task_def import DataFormat, EncoderModelType

logger = create_logger(
    __name__,
    to_disk=True,
    log_file='mt_dnn_feature_extractor.log')


def load_data(file):
    rows = []
    cnt = 0
    is_single_sentence = False
    with open(file, encoding="utf8") as f:
        for line in f:
            blocks = line.strip().split('|||')
            if len(blocks) == 2:
                sample = {
                    'uid': str(cnt),
                    'premise': blocks[0],
                    'hypothesis': blocks[1],
                    'label': 0}
            else:
                is_single_sentence = True
                sample = {'uid': str(cnt), 'premise': blocks[0], 'label': 0}
            rows.append(sample)
            cnt += 1
    return rows, is_single_sentence


def build_data(data, max_seq_len, is_train=True, tokenizer=None):
    """Build data of sentence pair tasks
    """
    rows = []
    for idx, sample in enumerate(data):
        ids = sample['uid']
        premise = tokenizer.tokenize(sample['premise'])
        hypothesis = tokenizer.tokenize(sample['hypothesis'])
        label = sample['label']
        _truncate_seq_pair(premise, hypothesis, max_seq_len - 3)
        input_ids = tokenizer.convert_tokens_to_ids(
            ['[CLS]'] + hypothesis + ['[SEP]'] + premise + ['[SEP]'])
        type_ids = [0] * (len(hypothesis) + 2) + [1] * (len(premise) + 1)
        features = {
            'uid': ids,
            'label': label,
            'token_id': input_ids,
            'type_id': type_ids,
            'tokens': ['[CLS]'] + hypothesis + ['[SEP]'] + premise + ['[SEP]']}
        rows.append(features)
    return rows


def build_data_single(data, max_seq_len, tokenizer=None):
    """Build data of single sentence tasks
    """
    rows = []
    for idx, sample in enumerate(data):
        ids = sample['uid']
        premise = tokenizer.tokenize(sample['premise'])
        label = sample['label']
        if len(premise) > max_seq_len - 3:
            premise = premise[:max_seq_len - 3]
        input_ids = tokenizer.convert_tokens_to_ids(
            ['[CLS]'] + premise + ['[SEP]'])
        type_ids = [0] * (len(premise) + 2)
        features = {
            'uid': ids,
            'label': label,
            'token_id': input_ids,
            'type_id': type_ids,
            'tokens': ['[CLS]'] + premise + ['[SEP]']}
        rows.append(features)
    return rows


def model_config(parser):
    parser.add_argument('--update_bert_opt', default=0, type=int)
    parser.add_argument('--multi_gpu_on', action='store_true')
    parser.add_argument('--mem_cum_type', type=str, default='simple',
                        help='bilinear/simple/defualt')
    parser.add_argument('--answer_num_turn', type=int, default=5)
    parser.add_argument('--answer_mem_drop_p', type=float, default=0.1)
    parser.add_argument('--answer_att_hidden_size', type=int, default=128)
    parser.add_argument('--answer_att_type', type=str, default='bilinear',
                        help='bilinear/simple/defualt')
    parser.add_argument('--answer_rnn_type', type=str, default='gru',
                        help='rnn/gru/lstm')
    parser.add_argument('--answer_sum_att_type', type=str, default='bilinear',
                        help='bilinear/simple/defualt')
    parser.add_argument('--answer_merge_opt', type=int, default=1)
    parser.add_argument('--answer_mem_type', type=int, default=1)
    parser.add_argument('--answer_dropout_p', type=float, default=0.1)
    parser.add_argument('--answer_weight_norm_on', action='store_true')
    parser.add_argument('--dump_state_on', action='store_true')
    parser.add_argument('--answer_opt', type=int, default=0, help='0,1')
    parser.add_argument('--label_size', type=str, default='3')
    parser.add_argument('--mtl_opt', type=int, default=0)
    parser.add_argument('--ratio', type=float, default=0)
    parser.add_argument('--mix_opt', type=int, default=0)
    parser.add_argument('--init_ratio', type=float, default=1)
    return parser


def train_config(parser):
    parser.add_argument('--cuda', type=bool, default=torch.cuda.is_available(),
                        help='whether to use GPU acceleration.')
    parser.add_argument('--optimizer', default='adamax',
                        help='supported optimizer: adamax, sgd, adadelta, adam')
    parser.add_argument('--grad_clipping', type=float, default=0)
    parser.add_argument('--global_grad_clipping', type=float, default=1.0)
    parser.add_argument('--weight_decay', type=float, default=0)
    parser.add_argument('--learning_rate', type=float, default=5e-5)
    parser.add_argument('--momentum', type=float, default=0)
    parser.add_argument('--warmup', type=float, default=0.1)
    parser.add_argument('--warmup_schedule', type=str, default='warmup_linear')
    parser.add_argument('--vb_dropout', action='store_false')
    parser.add_argument('--dropout_p', type=float, default=0.1)
    parser.add_argument('--dropout_w', type=float, default=0.000)
    parser.add_argument('--bert_dropout_p', type=float, default=0.1)
    parser.add_argument('--ema_opt', type=int, default=0)
    parser.add_argument('--ema_gamma', type=float, default=0.995)
    # scheduler
    parser.add_argument('--have_lr_scheduler', dest='have_lr_scheduler', action='store_false')
    parser.add_argument('--multi_step_lr', type=str, default='10,20,30')
    parser.add_argument('--freeze_layers', type=int, default=-1)
    parser.add_argument('--embedding_opt', type=int, default=0)
    parser.add_argument('--lr_gamma', type=float, default=0.5)
    parser.add_argument('--bert_l2norm', type=float, default=0.0)
    parser.add_argument('--scheduler_type', type=str, default='ms', help='ms/rop/exp')
    parser.add_argument('--output_dir', default='checkpoint')
    parser.add_argument('--seed', type=int, default=2018, 
                        help='random seed for data shuffling, embedding init, etc.')
    parser.add_argument('--encoder_type', type=int, default=EncoderModelType.BERT)
    #fp 16
    parser.add_argument('--fp16', action='store_true',
                        help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit")
    parser.add_argument('--fp16_opt_level', type=str, default='O1',
                        help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
                             "See details at https://nvidia.github.io/apex/amp.html")
    return parser


def set_config(parser):
    parser.add_argument("--finput", default=None, type=str, required=True)
    parser.add_argument("--foutput", default=None, type=str, required=True)
    parser.add_argument("--bert_model", default=None, type=str, required=True,
        help='Bert model: bert-base-uncased')
    parser.add_argument( "--checkpoint", default=None, type=str, required=True,
        help='model parameters')
    parser.add_argument( "--do_lower_case", action='store_true',
        help="Set this flag if you are using an uncased model.")
    parser.add_argument("--layers", default="10,11", type=str)
    parser.add_argument("--max_seq_length", default=512, type=int, help='')
    parser.add_argument("--batch_size", default=4, type=int)


def process_data(args):
    tokenizer = BertTokenizer.from_pretrained(
        args.bert_model, do_lower_case=args.do_lower_case)
    path = args.finput
    data, is_single_sentence = load_data(path)
    if is_single_sentence:
        tokened_data = build_data_single(
            data, max_seq_len=args.max_seq_length, tokenizer=tokenizer)
    else:
        tokened_data = build_data(
            data,
            max_seq_len=args.max_seq_length,
            tokenizer=tokenizer)
    return tokened_data, is_single_sentence

def dump_data(data, path):
    with open(path, 'w', encoding='utf-8') as writer:
        for sample in data:
            writer.write('{}\n'.format(json.dumps(sample)))

def main():
    parser = argparse.ArgumentParser()
    model_config(parser)
    set_config(parser)
    train_config(parser)
    args = parser.parse_args()
    encoder_type = args.encoder_type
    layer_indexes = [int(x) for x in args.layers.split(",")]
    set_environment(args.seed)
    # process data
    data, is_single_sentence = process_data(args)
    data_type = DataFormat.PremiseOnly if is_single_sentence else DataFormat.PremiseAndOneHypothesis
    fout_temp = '{}.tmp'.format(args.finput)
    dump_data(data, fout_temp)
    collater = Collater(is_train=False, encoder_type=encoder_type)
    dataset = SingleTaskDataset(fout_temp, False, maxlen=args.max_seq_length, data_type=data_type)
    batcher = DataLoader(dataset, batch_size=args.batch_size, collate_fn=collater.collate_fn, pin_memory=args.cuda)
    opt = vars(args)
    # load model
    if os.path.exists(args.checkpoint):
        state_dict = torch.load(args.checkpoint)
        config = state_dict['config']
        config['dump_feature'] = True
        opt.update(config)
    else:
        logger.error('#' * 20)
        logger.error(
            'Could not find the init model!\n The parameters will be initialized randomly!')
        logger.error('#' * 20)
        return
    num_all_batches = len(batcher)
    model = MTDNNModel(
        opt,
        state_dict=state_dict,
        num_train_step=num_all_batches)
    if args.cuda:
        model.cuda()

    features_dict = {}
    for batch_meta, batch_data in batcher:
        batch_meta, batch_data = Collater.patch_data(args.cuda, batch_meta, batch_data)
        all_encoder_layers, _ = model.extract(batch_meta, batch_data)
        embeddings = [all_encoder_layers[idx].detach().cpu().numpy()
                      for idx in layer_indexes]

        #import pdb; pdb.set_trace()
        uids = batch_meta['uids']
        masks = batch_data[batch_meta['mask']].detach().cpu().numpy().tolist()
        for idx, uid in enumerate(uids):
            slen = sum(masks[idx])
            features = {}
            for yidx, layer in enumerate(layer_indexes):
                features[layer] = str(embeddings[yidx][idx][:slen].tolist())
            features_dict[uid] = features

    # save features
    with open(args.foutput, 'w', encoding='utf-8') as writer:
        for sample in data:
            uid = sample['uid']
            tokens = sample['tokens']
            feature = features_dict[uid]
            feature['tokens'] = tokens
            feature['uid'] = uid
            writer.write('{}\n'.format(json.dumps(feature)))


if __name__ == "__main__":
    main()