Upload representation.py

representation.py

@@ -0,0 +1,90 @@
+import torch
+import torch.nn as nn
+from transformers import AutoModel, AutoTokenizer
+
+DEFAULT_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class TransformerRepresentation(nn.Module):
+    def __init__(self, model_name='bert-base-uncased',
+                 transformer_kwargs={'attention_probs_dropout_prob': 0.1,
+                                     'hidden_dropout_prob': 0.1},
+                 device=DEFAULT_DEVICE):
+        super(TransformerRepresentation, self).__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModel.from_pretrained(model_name,
+                                               output_hidden_states=True,
+                                               **transformer_kwargs)
+        self.embedding_dim = self.model.config.hidden_size
+        self.device = device
+
+    @staticmethod
+    def add_subword_maps(texts, encodings):
+        for encoding, t in zip(encodings, texts):
+            encoding.subword_map = [encoding.word_to_tokens(i)
+                                    for i, _ in enumerate(t)]
+
+    @staticmethod
+    def apply_token_pooling_strategy(outputs, encodings, strategy='first'):
+        """
+        Applies a token pooling strategy for pretokenized inputs based on
+        a sub-word mapping of words to tokens.
+
+        :param outputs: Output of the application of a `TransformerRepresentation.model` to a pretokenized input.
+        :param encodings: Encodings from the application of `TransformerRepresentation.tokenizer` to a pretokenized input.
+        :param strategy: One of ['first', 'last', 'sum', 'average']. Defaults to 'first'.
+        :return:
+        """
+        vec_map = [[vecs[m[0]:m[1]] for m in encoding.subword_map
+                    if m is not None]  # Only return vectors for words that were not truncated during tokenization
+                   for vecs, encoding
+                   in zip(outputs.last_hidden_state.unbind(), encodings)]
+        if strategy == 'first':
+            return [torch.stack([vec[0] for vec in vm]) if vm else torch.zeros(0)  for vm in vec_map]
+        elif strategy == 'last':
+            return [torch.stack([vec[-1] for vec in vm]) if vm else torch.zeros(0)  for vm in vec_map]
+        elif strategy == 'sum':
+            return [torch.stack([torch.sum(vec, dim=0) for vec in vm]) if vm else torch.zeros(0) for vm in vec_map]
+        elif strategy == 'average':
+            return [torch.stack([torch.sum(vec, dim=0)/len(vec) for vec in vm]) if vm else torch.zeros(0)  for vm in vec_map]
+        return vec_map
+
+    def add_special_tokens(self, tokens):
+        self.tokenizer.add_special_tokens({'additional_special_tokens': self.tokenizer.additional_special_tokens + tokens})
+        self.model.resize_token_embeddings(len(self.tokenizer))
+
+    def forward(self, text, is_pretokenized=False, add_special_tokens=True, token_pooling='first'):
+        inputs = self.tokenizer(text, padding='longest',
+                                is_split_into_words=is_pretokenized,
+                                add_special_tokens=add_special_tokens,
+                                return_tensors='pt',
+                                max_length=512,
+                                truncation=True).to(self.device)
+        output = self.model(**inputs.to(self.device))
+        if is_pretokenized:
+            self.add_subword_maps(text, [i for i in inputs.encodings])
+            output.pooled_tokens = self.apply_token_pooling_strategy(
+                output, [i for i in inputs.encodings], strategy=token_pooling)
+        return output
+
+
+if __name__ == 'main':
+    toks = ['Tom', 'Thabane', 'resigned', 'in', 'October', 'last', 'year',
+            'to', 'form', 'the', 'All', 'Basotho', 'Convention', '-LRB-',
+            'ABC', '-RRB-', ',', 'crossing', 'the', 'floor', 'with', '17',
+            'members', 'of', 'parliament', ',', 'causing', 'constitutional',
+            'monarch', 'King', 'Letsie', 'III', 'to', 'dissolve',
+            'parliament', 'and', 'call', 'the', 'snap', 'election', '.']
+    e1_type = 'PERSON'
+    e2_type = 'ORGANIZATION'
+    e1_tokens = [0, 1]
+    e2_tokens = [10, 12]
+    text = [['EU', 'rejects', 'German', 'call', 'to', 'boycott', 'British', 'lamb', '.'],
+            ['Peter', 'Blackburn'],
+            ['BRUSSELS', '1996-08-22'],
+            ['The', 'European', 'Commission', 'said', 'on', 'Thursday', 'it', 'disagreed', 'with', 'German', 'advice', 'to', 'consumers', 'to', 'shun', 'British', 'lamb', 'until', 'scientists', 'determine', 'whether', 'mad', 'cow', 'disease', 'can', 'be', 'transmitted', 'to', 'sheep', '.'],
+            ['Germany', "'s", 'representative', 'to', 'the', 'European', 'Union', "'s", 'veterinary', 'committee', 'Werner', 'Zwingmann', 'said', 'on', 'Wednesday', 'consumers', 'should', 'buy', 'sheepmeat', 'from', 'countries', 'other', 'than', 'Britain', 'until', 'the', 'scientific', 'advice', 'was', 'clearer', '.'],
+            ['"', 'We', 'do', "n't", 'support', 'any', 'such', 'recommendation', 'because', 'we', 'do', "n't", 'see', 'any', 'grounds', 'for', 'it', ',', '"', 'the', 'Commission', "'s", 'chief', 'spokesman', 'Nikolaus', 'van', 'der', 'Pas', 'told', 'a', 'news', 'briefing', '.'],
+            ['He', 'said', 'further', 'scientific', 'study', 'was', 'required', 'and', 'if', 'it', 'was', 'found', 'that', 'action', 'was', 'needed', 'it', 'should', 'be', 'taken', 'by', 'the', 'European', 'Union', '.'],
+            ['He', 'said', 'a', 'proposal', 'last', 'month', 'by', 'EU', 'Farm', 'Commissioner', 'Franz', 'Fischler', 'to', 'ban', 'sheep', 'brains', ',', 'spleens', 'and', 'spinal', 'cords', 'from', 'the', 'human', 'and', 'animal', 'food', 'chains', 'was', 'a', 'highly', 'specific', 'and', 'precautionary', 'move', 'to', 'protect', 'human', 'health', '.']]
+    model = TransformerRepresentation()