coverage_eval.py

"""
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 coverage(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()

    # path to dataset
    data_path = args.data_path

    # load dataset
    batcher = BartBatcher(tokenizer, model.config, [data_path], torch_device, synthetic=args.synthetic)

    scores = []
    while batcher.epoch_counter < 1:
        input_ids, att_mask, targ_ids, targ_att_mask = batcher.get_a_batch(batch_size=1)

        # 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

        if args.ctype == 'var':
            # 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 = np.var(norm_max)
        elif args.ctype == 'ent':
            # entropy
            entropies = entropy(CAM, axis=0) / len(CAM)
            # coverage score
            score = entropy(entropies) / len(entropies)

        scores.append(score)

    return (sum(scores) / len(scores))
        
        

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("-data_path", default="/synthetic_data", type=str, help="Path to data with text and QA pairs.")
    parser.add_argument("-ctype", default='var', choices=['var', 'ent'], help="Coverage metric to calculate: var or ent.")
    parser.add_argument("-synthetic", type=bool, help="True if the dataset is synthetic. None if dataset is not.")

    args = parser.parse_args()

    coverage_score = coverage(args)
    print("Coverage Socre: ",coverage_score)