simon-clmtd/exbert

1import numpy as np2 3SEP = '[SEP]'4CLS = '[CLS]'5MASK = '[MASK]'6 7def drop_bad_inds(arr, left_drop, right_drop):8    """Given the 4d array returned by attentions of shape (n_layer, n_head, n_left_text, n_right_text),9        return that array modified to drop ind1 from n_left_text and ind2 from n_right_text10    """11    # print("Length of left drop: ", len(left_drop))12    # print("Length of right drop: ", len(left_drop))13    print("Shape of arr: ", arr.shape)14    arr = arr[:, :, ~left_drop, :]15 16    # Keys and queries don't match in the final dimension17    if arr.shape[-1] == len(right_drop):18        arr = arr[:, :, :, ~right_drop]19 20    return arr21 22def strip_attention(attention):23    """Given an attention output of the BERT model, 24    return the same object without CLS and SEP token weightings25 26    NOTE: Not currently fixing key and query27    """28    attention_out = {}29 30    # Iterate through sentence combinations31    # Need queries, keys, att, left_text, right_text32    for i, (k, v) in enumerate(attention.items()):33        stripped_resp = {}34 35        left_tokens = np.array(v['left_text'])36        right_tokens = np.array(v['right_text'])37        att = np.array(v['att'])38        # key = np.array(v['keys'])39        # quer = np.array(v['queries'])40 41        left_drop = (left_tokens == CLS) | (left_tokens == SEP)42        right_drop = (right_tokens == CLS) | (right_tokens == SEP)43        44        att_out = drop_bad_inds(att, left_drop, right_drop)45        # key_out = drop_bad_inds(key, left_drop, right_drop)46        # quer_out = drop_bad_inds(quer, left_drop, right_drop)47        left_out = left_tokens[~left_drop]48        right_out = right_tokens[~right_drop]49 50        # assert att_out.shape[:3] == key_out.shape[:3] == quer_out.shape[:3]51        assert att_out.shape[2] == len(left_out)52        assert att_out.shape[3] == len(right_out)53 54        stripped_resp['att'] = att_out.tolist()55        stripped_resp['keys'] = v['keys']56        stripped_resp['queries'] = v['queries']57        stripped_resp['left_text'] = left_out.tolist()58        stripped_resp['right_text'] = right_out.tolist()59 60        attention_out[k] = stripped_resp61 62    return attention_out63 64def mask_attention(deets, maskA, maskB):65    """Deets have form:66 67        tokens_a, tokens_b, query_tensor.data.numpy(), key_tensor.data.numpy(), attn_tensor.data.numpy()68 69    Take the first two in tuple and mask according to maskA and maskB which are lists of indices to mask70    """71 72    tokens_a = np.array(deets[0])73    tokens_a[maskA] = MASK74    tokens_a.tolist()75 76    tokens_b = np.array(deets[1])77    tokens_b[maskb] = MASK78    tokens_b.tolist()79 80    deets[0] = tokens_a.tolist()81    deets[1] = tokens_b.tolist()82 83    return deets