Simplify model_api in server

environment.yml

@@ -5,21 +5,14 @@ channels:
   - defaults
   - anaconda
 dependencies:
+  - pip>=19.0.3
   - python=3.7
   - connexion=1.5.3
-  - h5py
-  - spacy
   - boto3
   - regex
   - flask-cors
-  - faiss-cpu
   - jinja2=2.10
-  - numpy=1.16.2
-  - olefile=0.46
-  - pickleshare=0.7.5
-  - pillow=5.4.1
-  - pip=19.0.3
-  - pytorch=1.0.1
-  - sacremoses
-  - pip:
-    - sentencepiece
+  - numpy
+  - pytorch
+  - torchvision
+  - transformers

server/model_api.py

@@ -0,0 +1,161 @@
+from typing import List, Union, Tuple
+
+import torch
+from transformers import AutoConfig, AutoTokenizer, AutoModelWithLMHead, AutoModel
+
+from transformer_formatter import TransformerOutputFormatter
+from utils.f import delegates, pick, memoize
+
+def get_model_tok(mname):
+    conf = AutoConfig.from_pretrained(mname, output_attentions=True, output_past=False)
+    tok = AutoTokenizer.from_pretrained(mname, config=conf)
+    model = AutoModelWithLMHead.from_pretrained(mname, config=conf)
+    return model, tok
+
+class ModelDetails:
+    """Wraps a transformer model and tokenizer to prepare inputs to the frontend visualization"""
+    def __init__(self, mname):
+        self.mname = mname
+        self.model, self.tok = get_model_tok(self.mname)
+        self.model.eval()
+        self.config = self.model.config
+    
+    def from_sentence(self, sentence: str) -> TransformerOutputFormatter:
+        """Get attentions and word probabilities from a sentence. Special tokens are automatically added if a sentence is passed.
+        
+        Args:
+            sentence: The input sentence to tokenize and analyze. 
+        """
+        tokens = self.tok.tokenize(sentence)
+
+        return self.from_tokens(tokens, sentence, add_special_tokens=True)
+
+    def from_tokens(
+        self, tokens: List[str], orig_sentence:str, add_special_tokens:bool=False, mask_attentions:bool=False, topk:int=5
+    ) -> TransformerOutputFormatter:
+        """Get formatted attention and predictions from a list of tokens.
+
+        Args:
+            tokens: Tokens to analyze
+            orig_sentence: The sentence the tokens came from (needed to help organize the output)
+            add_special_tokens: Whether to add special tokens like CLS / <|endoftext|> to the tokens.
+                If False, assume the tokens already have the special tokens
+            mask_attentions: If True, do not pay attention to attention patterns to special tokens through the model.
+            topk: How many top predictions to report
+        """
+        ids = self.tok.convert_tokens_to_ids(tokens)
+
+        # For GPT2, add the beginning of sentence token to the input. Note that this will work on all models but XLM
+        bost = self.tok.bos_token_id
+        clst = self.tok.cls_token_id
+        if (bost is not None) and (bost != clst) and add_special_tokens:
+            ids.insert(0, bost)
+
+        inputs = self.tok.prepare_for_model(ids, add_special_tokens=add_special_tokens, return_tensors="pt")
+        parsed_input = self.parse_inputs(inputs, mask_attentions=mask_attentions)
+        output = self.model(parsed_input['input_ids'], attention_mask=parsed_input['attention_mask'])
+
+        logits, atts = self.choose_logits_att(output)
+        words, probs = self.logits2words(logits, topk)
+        tokens = self.view_ids(inputs["input_ids"])
+
+        formatted_output = TransformerOutputFormatter(
+            orig_sentence,
+            tokens,
+            inputs["special_tokens_mask"],
+            atts,
+            words,
+            probs.tolist(),
+            self.config
+        )
+
+        return formatted_output
+
+    def choose_logits_att(self, out:Tuple) -> Tuple:
+        """Select from the model's output the logits and the attentions, switching on model name
+        
+        Args:
+            out: Output from the model's forward pass
+
+        Returns:
+            (logits: tensor((bs, N)), attentions: Tuple[tensor(())])
+        """
+        if 't5' in self.mname:
+            logits, _, atts = out
+        else:
+            logits, atts = out
+
+        print("Logits: ", logits)
+        print("atts: ", atts[0].shape)
+        return logits, atts
+
+    def logits2words(self, logits, topk):
+        """Convert logit probabilities into words from the tokenizer's vocabulary.
+        
+        """
+        probs, idxs = torch.topk(torch.softmax(logits.squeeze(0), 1), topk)
+        words = [self.tok.convert_ids_to_tokens(i) for i in idxs]
+        return words, probs
+
+    def view_ids(self, ids: Union[List[int], torch.Tensor]) -> List[str]:
+        """View what the tokenizer thinks certain ids are for a single input"""
+        if type(ids) == torch.Tensor:
+            # Remove batch dimension
+            ids = ids.squeeze(0).tolist()
+
+        out = self.tok.convert_ids_to_tokens(ids)
+        return out
+
+    def parse_inputs(self, inputs, mask_attentions=False):
+        """Parse the output from `tokenizer.prepare_for_model` to the desired attention mask from special tokens
+
+        Args:
+            - inputs: The output of `tokenizer.prepare_for_model`.
+                A dict with keys: {'special_token_mask', 'token_type_ids', 'input_ids'}
+            - mask_attentions: Flag indicating whether to mask the attentions or not
+
+        Returns:
+            Dict with keys: {'input_ids', 'token_type_ids', 'attention_mask', 'special_tokens_mask'}
+
+        Usage:
+
+            ```
+            s = "test sentence"
+
+            # from raw sentence to tokens
+            tokens = tokenizer.tokenize(s)
+
+            # From tokens to ids
+            ids = tokenizer.convert_tokens_to_ids(tokens)
+
+            # From ids to input
+            inputs = tokenizer.prepare_for_model(ids, return_tensors='pt')
+
+            # Parse the input. Optionally mask the special tokens from the analysis.
+            parsed_input = parse_inputs(inputs)
+
+            # Run the model, pick from this output whatever inputs you want
+            from utils.f import pick
+            out = model(**pick(['input_ids'], parse_inputs(inputs)))
+            ```
+        """
+
+        out = inputs.copy()
+
+        # DEFINE SPECIAL TOKENS MASK
+        if "special_tokens_mask" not in inputs.keys():
+            special_tokens = set([self.tok.unk_token_id, self.tok.cls_token_id, self.tok.sep_token_id, self.tok.bos_token_id, self.tok.eos_token_id, self.tok.pad_token_id])
+            in_ids = inputs['input_ids'][0]
+            special_tok_mask = [1 if int(i) in special_tokens else 0 for i in in_ids]
+            inputs['special_tokens_mask'] = special_tok_mask
+
+        if mask_attentions:
+            out["attention_mask"] = torch.tensor(
+                [int(not i) for i in inputs.get("special_tokens_mask")]
+            ).unsqueeze(0)
+        else:
+            out["attention_mask"] = torch.tensor(
+                [1 for i in inputs.get("special_tokens_mask")]
+            ).unsqueeze(0)
+
+        return out

server/setup.py

@@ -3,8 +3,8 @@ from setuptools import setup, find_packages
 requires = [] # Let conda handle requires
 
 setup(
-    name="exbert",
-    description="Vis",
+    name="exformer",
+    description="Just the attention vis of exbert",
     packages=find_packages(),
     author="IBM Research AI",
     include_package_data=True,

server/transformer_formatter.py

@@ -4,7 +4,6 @@ import numpy as np
 import torch
 import json
 
-from spacyface.simple_spacy_token import SimpleSpacyToken
 from utils.token_processing import fix_byte_spaces
 from utils.gen_utils import map_nlist
 
@@ -14,8 +13,8 @@ def round_return_value(attentions, ndigits=5):
     
     attentions: {
         'aa': {
-            left.embeddings & contexts
-            right.embeddings & contexts
+            left
+            right
             att
         }
     }
@@ -25,19 +24,6 @@ def round_return_value(attentions, ndigits=5):
     nested_rounder = partial(map_nlist, rounder)
     new_out = attentions  # Modify values to save memory
     new_out["aa"]["att"] = nested_rounder(attentions["aa"]["att"])
-    new_out["aa"]["left"]["embeddings"] = nested_rounder(
-        attentions["aa"]["left"]["embeddings"]
-    )
-    new_out["aa"]["left"]["contexts"] = nested_rounder(
-        attentions["aa"]["left"]["contexts"]
-    )
-
-    new_out["aa"]["right"]["embeddings"] = nested_rounder(
-        attentions["aa"]["right"]["embeddings"]
-    )
-    new_out["aa"]["right"]["contexts"] = nested_rounder(
-        attentions["aa"]["right"]["contexts"]
-    )
 
     return new_out
 
@@ -60,71 +46,40 @@ class TransformerOutputFormatter:
     def __init__(
         self,
         sentence: str,
-        tokens: List[SimpleSpacyToken],
+        tokens: List[str],
         special_tokens_mask: List[int],
         att: Tuple[torch.Tensor], 
-        embeddings: Tuple[torch.Tensor],
-        contexts: Tuple[torch.Tensor],
         topk_words: List[List[str]],
-        topk_probs: List[List[float]]
+        topk_probs: List[List[float]],
+        model_config
     ):
         assert len(tokens) > 0, "Cannot have an empty token output!"
 
-        modified_embeddings = flatten_batch(embeddings)
         modified_att = flatten_batch(att)
-        modified_contexts = flatten_batch(contexts)
 
         self.sentence = sentence
         self.tokens = tokens
         self.special_tokens_mask = special_tokens_mask
-        self.embeddings = modified_embeddings
         self.attentions = modified_att
-        self.raw_contexts = modified_contexts
         self.topk_words = topk_words
         self.topk_probs = topk_probs
+        self.model_config = model_config
 
-        self.n_layers = len(contexts) # With +1 for buffer layer at the beginning
-        _, self.__len, self.n_heads, self.hidden_dim = contexts[0].shape
-
-    @property
-    def contexts(self):
-        """Combine the head and the context dimension as it is passed forward in the model"""
-        return squeeze_contexts(self.raw_contexts)
-
-    @property
-    def normed_embeddings(self):
-        ens = tuple([torch.norm(e, dim=-1) for e in self.embeddings])
-        normed_es = tuple([e / en.unsqueeze(-1) for e, en in zip(self.embeddings, ens)])
-        return normed_es
-
-    @property
-    def normed_contexts(self):
-        """Normalize each by head"""
-        cs = self.raw_contexts
-        cns = tuple([torch.norm(c, dim=-1) for c in cs])
-        normed_cs = tuple([c / cn.unsqueeze(-1) for c, cn in zip(cs, cns)])
-        squeezed_normed_cs = squeeze_contexts(normed_cs)
-        return squeezed_normed_cs
+        self.n_layer = self.model_config.n_layer
+        self.n_head = self.model_config.n_head
+        self.hidden_dim = self.model_config.n_embd
+
+        self.__len = len(tokens)# Get the number of tokens in the input
+        assert self.__len == self.attentions[0].shape[-1], "Attentions don't represent the passed tokens!"
     
     def to_json(self, layer:int, ndigits=5):
         """The original API expects the following response:
 
         aa: {
             att: number[][][]
-            left: <FullSingleTokenInfo[]>
-            right: <FullSingleTokenInfo[]>
+            left: List[str]
+            right: List[str]
         }
-
-        FullSingleTokenInfo:
-            {
-                text: string
-                embeddings: number[]
-                contexts: number[]
-                bpe_token: string
-                bpe_pos: string
-                bpe_dep: string
-                bpe_is_ent: boolean
-            }
         """
         # Convert the embeddings, attentions, and contexts into list. Perform rounding
 
@@ -133,25 +88,16 @@ class TransformerOutputFormatter:
 
         def tolist(tens): return [t.tolist() for t in tens]
 
-        def to_resp(tok: SimpleSpacyToken, embeddings: List[float], contexts: List[float], topk_words, topk_probs):
+        def to_resp(tok: str, topk_words, topk_probs):
             return {
-                "text": tok.token,
-                "bpe_token": tok.token,
-                "bpe_pos": tok.pos,
-                "bpe_dep": tok.dep,
-                "bpe_is_ent": tok.is_ent,
-                "embeddings": nested_rounder(embeddings),
-                "contexts": nested_rounder(contexts),
+                "text": tok,
                 "topk_words": topk_words,
                 "topk_probs": nested_rounder(topk_probs)
             }
 
-        side_info = [to_resp(t, e, c, w, p) for t,e,c,w,p in zip(
-                                                                self.tokens, 
-                                                                tolist(self.embeddings[layer]), 
-                                                                tolist(self.contexts[layer]),
-                                                                self.topk_words,
-                                                                self.topk_probs)]
+        side_info = [to_resp(t, w, p) for t,w,p in zip( self.tokens, 
+                                                        self.topk_words,
+                                                        self.topk_probs)]
 
         out = {"aa": {
             "att": nested_rounder(tolist(self.attentions[layer])),
@@ -164,42 +110,6 @@ class TransformerOutputFormatter:
     def display_tokens(self, tokens):
         return fix_byte_spaces(tokens)
 
-    def to_hdf5_meta(self):
-        """Output metadata information to store as hdf5 metadata for a group"""
-        token_dtype = self.tokens[0].hdf5_token_dtype
-        out = {k: np.array([t[k] for t in self.tokens], dtype=np.dtype(dtype)) for k, dtype in token_dtype}
-        out['sentence'] = self.sentence
-        return out
-
-    def to_hdf5_content(self, do_norm=True):
-        """Return dictionary of {attentions, embeddings, contexts} formatted as array for hdf5 file"""
-
-        def get_embeds(c): 
-            if do_norm: return c.normed_embeddings
-            return c.embeddings
-
-        def get_contexts(c):
-            if do_norm: return c.normed_contexts
-            return c.contexts
-
-        embeddings = to_numpy(get_embeds(self))
-        contexts = to_numpy(get_contexts(self))
-        atts = to_numpy(self.attentions)
-
-        return {
-            "embeddings": embeddings,
-            "contexts": contexts,
-            "attentions": atts
-        }
-
-    @property
-    def searchable_embeddings(self):
-        return np.array(list(map(to_searchable, self.embeddings)))
-
-    @property
-    def searchable_contexts(self):
-        return np.array(list(map(to_searchable, self.contexts)))
-
     def __repr__(self):
         lim = 50
         if len(self.sentence) > lim: s = self.sentence[:lim - 3] + "..."

server/utils/gen_utils.py

@@ -1,15 +1,8 @@
-import spacy
 from copy import deepcopy
 import numpy as np
 from functools import partial
 from .f import memoize
 
-def add_base_exceptions(language_exceptions): 
-    merged = {}
-    merged.update(language_exceptions)
-    merged.update(spacy.lang.tokenizer_exceptions.BASE_EXCEPTIONS)
-    return merged
-
 def check_key_len(d, length):
     for k, v in d.items():
         if len(v) != length:

server/utils/token_processing.py

@@ -1,191 +1,16 @@
-"""Defines the important metadata to extract for each token.
-
-If adding more metadata, modify the definitions in `to_spacy_meta` and `meta_to_hdf5`
-"""
-import h5py
 import numpy as np
-import spacy
 from transformers.tokenization_bert import BertTokenizer
 from .f import flatten_, assoc, memoize, GetAttr
 
 from typing import List
 
 def fix_byte_spaces(toks: List[str]) -> List[str]:
-        return [t.replace("\u0120", " ") for t in toks]
-
-# NOTE: If you want to change anything that is extracted from the SPACY token, change the functions below.
-# ====================================================================================================
-def simplify_spacy_token(t):
-    """Extract important information from spacy token into a simple dictionary"""
-    def check_ent(tok):
-        OUT_OF_ENT = 2
-        NO_ENT_DEFINED = 0
-        return tok.ent_iob != OUT_OF_ENT and tok.ent_iob != NO_ENT_DEFINED
-
-    return {
-        "token": t.text,
-        "pos": t.pos_,
-        "dep": t.dep_,
-        "norm": t.norm_,
-        "tag": t.tag_,
-        "lemma": t.lemma_,
-        "head": t.head,
-        "is_ent": check_ent(t),
-    }
-
-def null_token_filler(token_text):
-    return {
-        "token": token_text,
-        "pos": None,
-        "dep": None,
-        "norm": None,
-        "tag": None,
-        "lemma": None,
-        "head": None,
-        "is_ent": None,
-    }
-
-token_dtype = [
-        ("token", h5py.special_dtype(vlen=str)),
-        ("pos", h5py.special_dtype(vlen=str)),
-        ("dep", h5py.special_dtype(vlen=str)),
-        ("norm", h5py.special_dtype(vlen=str)),
-        ("tag", h5py.special_dtype(vlen=str)),
-        ("lemma", h5py.special_dtype(vlen=str)),
-        ("head", h5py.special_dtype(vlen=str)),
-        ("is_ent", np.bool_),
-    ]
-# ====================================================================================================
+        return [t.replace("\u0120", " ").replace("\u010A", "\\n") for t in toks]
 
 @memoize
 def get_bpe(bpe_pretrained_name_or_path):
     return BertTokenizer.from_pretrained(bpe_pretrained_name_or_path)
 
-@memoize
-def get_spacy(spacy_name):
-    return spacy.load(spacy_name)
-
-class TokenAligner:
-    def __init__(
-        self,
-        bpe_pretrained_name_or_path="bert-base-uncased",
-        spacy_name="en_core_web_sm",
-    ):
-        """Create a wrapper around a sentence such that the spacy and BPE tokens can be aligned"""
-        self.bpe = get_bpe(bpe_pretrained_name_or_path)
-        self.nlp = get_spacy(spacy_name)
-
-    def fix_sentence(self, s):
-        return " ".join(self.to_spacy(s))
-
-    def to_spacy(self, s):
-        """Convert a sentence to spacy tokens. 
-        
-        Note that all contractions are removed in lieu of the word they shorten by taking the 'norm' of the word as defined by spacy.
-        """
-        doc = self.nlp(s)
-        tokens = [t.norm_ for t in doc]
-        return tokens
-
-    def to_spacy_text(self, s):
-        """Convert a sentence into the raw tokens as spacy would.
-        
-        No contraction expansion."""
-        doc = self.nlp(s)
-        tokens = [t.text for t in doc]
-        return tokens
-
-    def to_bpe(self, s):
-        """Convert a sentence to bpe tokens"""
-        s = self.fix_sentence(s)
-        s = self.to_bpe_text(s)
-        return s
-
-    def to_bpe_text(self, s):
-        """Convert a sentence to bpe tokens"""
-        return self.bpe.tokenize(s)
-
-    def to_spacy_meta(self, s):
-        """Convert a sentence to spacy tokens with important metadata"""
-        doc = self.nlp(s)
-        out = [simplify_spacy_token(t) for t in doc]
-        return out
-
-    def meta_to_hdf5(self, meta):
-        out_dtype = np.dtype(token_dtype)
-
-        out = [tuple([m[d[0]] for d in token_dtype]) for m in meta]
-        return np.array(out, dtype=out_dtype)
-
-    def meta_hdf5_to_obj(self, meta_hdf5):
-        assert len(meta_hdf5) != 0
-
-        keys = meta_hdf5[0].dtype.names
-        out = {k: [] for k in keys}
-
-        for m in meta_hdf5:
-            for k in m.dtype.names:
-                out[k].append(m[k])
-        return out
-
-    def to_spacy_hdf5(self, s):
-        """Get values for hdf5 store, each row being a tuple of the information desired"""
-        meta = self.to_spacy_meta(s)
-        return self.meta_to_hdf5(meta)
-
-    def to_spacy_hdf5_by_col(self, s):
-        """Get values for hdf5 store, organized as a dictionary into the metadata"""
-        h5_info = self.to_spacy_hdf5(s)
-        return self.meta_hdf5_to_obj(h5_info)
-
-    def bpe_from_meta_single(self, meta_token):
-        """Split a single spacy token with metadata into bpe tokens"""
-
-        bpe_tokens = self.to_bpe(meta_token["norm"])
-
-        # print(bpe_tokens)
-        return [assoc("token", b, meta_token) for b in bpe_tokens]
-
-    def bpe_from_spacy_meta(self, spacy_meta):
-        out = flatten_([self.bpe_from_meta_single(sm) for sm in spacy_meta])
-        return out
-
-    def to_bpe_meta(self, s):
-        """Convert a sentence to bpe tokens with metadata
-        
-        Removes all known contractions from input sentence `s`
-        """
-        bpe = self.to_bpe(s)
-        spacy_meta = self.to_spacy_meta(s)
-        return self.bpe_from_spacy_meta(spacy_meta)
-
-    def to_bpe_meta_from_tokens(self, sentence, bpe_tokens):
-        """Get the normal BPE metadata, and add nulls wherever a special_token appears"""
-        bpe_meta = self.to_bpe_meta(sentence)
-
-        new_bpe_meta = []
-        j = 0
-        for i, b in enumerate(bpe_tokens):
-            if b in self.bpe.all_special_tokens:
-                new_bpe_meta.append(null_token_filler(b))
-            else:
-                new_bpe_meta.append(bpe_meta[j])
-                j += 1
-
-        return new_bpe_meta
-
-    def to_bpe_hdf5(self, s):
-        """Format the metadata of a BPE tokenized setence into hdf5 format"""
-        meta = self.to_bpe_meta(s)
-        return self.meta_to_hdf5(meta)
-
-    def to_bpe_hdf5_by_col(self, s):
-        h5_info = self.to_bpe_hdf5(s)
-        return self.meta_hdf5_to_obj(h5_info)
-
-    def meta_tokenize(self, s):
-        return self.to_bpe_meta(s)
-
 # [String] -> [String]
 def remove_CLS_SEP(toks):
     return [t for t in toks if t not in set(["[CLS]", "[SEP]"])]