simon-clmtd/exbert

1from typing import List, Iterable, Tuple2from functools import partial3import numpy as np4import torch5import json6 7from utils.token_processing import fix_byte_spaces8from utils.gen_utils import map_nlist9 10 11def round_return_value(attentions, ndigits=5):12    """Rounding must happen right before it's passed back to the frontend because there is a little numerical error that's introduced converting back to lists13    14    attentions: {15        'aa': {16            left17            right18            att19        }20    }21    22    """23    rounder = partial(round, ndigits=ndigits)24    nested_rounder = partial(map_nlist, rounder)25    new_out = attentions  # Modify values to save memory26    new_out["aa"]["att"] = nested_rounder(attentions["aa"]["att"])27 28    return new_out29 30def flatten_batch(x: Tuple[torch.Tensor]) -> Tuple[torch.Tensor]:31    """Remove the batch dimension of every tensor inside the Iterable container `x`"""32    return tuple([x_.squeeze(0) for x_ in x])33 34def squeeze_contexts(x: Tuple[torch.Tensor]) -> Tuple[torch.Tensor]:35    """Combine the last two dimensions of the context."""36    shape = x[0].shape37    new_shape = shape[:-2] + (-1,)38    return tuple([x_.view(new_shape) for x_ in x])39 40def add_blank(xs: Tuple[torch.tensor]) -> Tuple[torch.Tensor]:41    """The embeddings have n_layers + 1, indicating the final output embedding."""42 43    return (torch.zeros_like(xs[0]),) + xs44 45class TransformerOutputFormatter:46    def __init__(47        self,48        sentence: str,49        tokens: List[str],50        special_tokens_mask: List[int],51        att: Tuple[torch.Tensor], 52        topk_words: List[List[str]],53        topk_probs: List[List[float]],54        model_config55    ):56        assert len(tokens) > 0, "Cannot have an empty token output!"57 58        modified_att = flatten_batch(att)59 60        self.sentence = sentence61        self.tokens = tokens62        self.special_tokens_mask = special_tokens_mask63        self.attentions = modified_att64        self.topk_words = topk_words65        self.topk_probs = topk_probs66        self.model_config = model_config67 68        try: 69            # GPT vals70            self.n_layer = self.model_config.n_layer71            self.n_head = self.model_config.n_head72            self.hidden_dim = self.model_config.n_embd73        except AttributeError:74            try: 75                # BERT vals76                self.n_layer = self.model_config.num_hidden_layers77                self.n_head = self.model_config.num_attention_heads78                self.hidden_dim = self.model_config.hidden_size79            except AttributeError: raise80 81 82        self.__len = len(tokens)# Get the number of tokens in the input83        assert self.__len == self.attentions[0].shape[-1], "Attentions don't represent the passed tokens!"84    85    def to_json(self, layer:int, ndigits=5):86        """The original API expects the following response:87 88        aa: {89            att: number[][][]90            left: List[str]91            right: List[str]92        }93        """94        # Convert the embeddings, attentions, and contexts into list. Perform rounding95 96        rounder = partial(round, ndigits=ndigits)97        nested_rounder = partial(map_nlist, rounder)98 99        def tolist(tens): return [t.tolist() for t in tens]100 101        def to_resp(tok: str, topk_words, topk_probs):102            return {103                "text": tok,104                "topk_words": topk_words,105                "topk_probs": nested_rounder(topk_probs)106            }107 108        side_info = [to_resp(t, w, p) for t,w,p in zip( self.tokens, 109                                                        self.topk_words,110                                                        self.topk_probs)]111 112        out = {"aa": {113            "att": nested_rounder(tolist(self.attentions[layer])),114            "left": side_info,115            "right": side_info116        }}117 118        return out119 120    def display_tokens(self, tokens):121        return fix_byte_spaces(tokens)122 123    def __repr__(self):124        lim = 50125        if len(self.sentence) > lim: s = self.sentence[:lim - 3] + "..."126        else: s = self.sentence[:lim]127 128        return f"TransformerOutput({s})"129 130    def __len__(self):131        return self.__len132        133def to_numpy(x): 134    """Embeddings, contexts, and attentions are stored as torch.Tensors in a tuple. Convert this to a numpy array135    for storage in hdf5"""136    return np.array([x_.detach().numpy() for x_ in x])137 138def to_searchable(t: Tuple[torch.Tensor]):139    return t.detach().numpy().astype(np.float32)