aws-dialogue-qa-pair-generation/get_coverage_scores.py at main · amazon-science/aws-dialogue-qa-pair-generation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
"""
This file calculates the coverage for a given dataset containing texts and corresponding QA Pairs.
"""

import sys
sys.path.append('mnt/efs/project')

import argparse
import json
import torch
import numpy as np
from scipy.stats import entropy

# bart model
from transformers import BartTokenizer, BartForConditionalGeneration

# dataloader
from loaders.dataloader_pt import BartBatcher

# gpu
torch_device = 'cuda' if torch.cuda.is_available() else 'cpu'

# entropy
from scipy.stats import entropy

def add_coverage_scores(args):

    #bart model
    tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-xsum")
    model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-xsum")

    model_config = model.config
    if torch_device=='cuda': model.cuda()

    with open(args.data_path, 'r') as f_in, open(args.output_path, 'w') as f_out:
        for line in f_in:
            ex = json.loads(line)
            dialogue = ' '.join( ex["dialogue"])
            question = ex["question"]
            answer = ex["answer"]
            qa_pair = "<q> " + question + " <a> " + answer

            ins = [ dialogue ]
            tar = [ qa_pair ]

            batch_encoded_inputs = tokenizer.batch_encode_plus(ins,
                                                               max_length=1024,
                                                               add_special_tokens=True,
                                                               padding=True,
                                                               truncation=True,
                                                               return_tensors='pt')
            input_ids = batch_encoded_inputs['input_ids']
            att_mask = batch_encoded_inputs['attention_mask']

            batch_encoded_targets = tokenizer.batch_encode_plus(tar,
                                                               add_special_tokens=True,
                                                               padding=True,
                                                               truncation=True,
                                                               return_tensors='pt')
            targ_ids = batch_encoded_targets['input_ids']
            targ_att_mask = batch_encoded_targets['attention_mask']

            if torch_device=='cuda':
                input_ids = input_ids.to(torch_device)
                att_mask = att_mask.to(torch_device)
                targ_ids = targ_ids.to(torch_device)
                targ_att_mask = targ_att_mask.to(torch_device)


            # bart to get CAM - Cross Attention Matrix
            with torch.no_grad():
                x = model(
                    input_ids = input_ids,
                    attention_mask = att_mask,
                    decoder_input_ids = targ_ids,
                    decoder_attention_mask = targ_att_mask,
                    output_attentions = True,
                )

            # average across layers to get CAM
            layers = []
            for layer in x.cross_attentions:
                layers.append( layer.detach().cpu().numpy() )
            CAM = np.mean(np.mean(np.mean(layers, axis=0), axis=0), axis=0)
            CAM = CAM[:, 1:-2] # drop start and end tokens

            # get max score along dialogue tokens
            maxs = np.amax(CAM, axis=0)
            # normalize
            norm = np.linalg.norm(maxs)
            norm_max = maxs / norm
            # coverage score
            score_var = np.var(norm_max)

            # entropy
            entropies = entropy(CAM, axis=0) / len(CAM)
            # coverage score
            score_ent = entropy(entropies) / len(entropies)

            ex["var_cov"] = str(score_var)
            ex["ent_cov"] = str(score_ent)

            json.dump(ex, f_out)
            f_out.write('\n')


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("-data_path", default="/Data", type=str, help="Path to data with text and QA pairs.")
    parser.add_argument("-output_path", default="/Data/Scores", type=str, help="Path to output file.")


    args = parser.parse_args()

    add_coverage_scores(args)