File size: 12,448 Bytes
6124176
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
import heapq
import itertools
from abc import ABC, abstractmethod
from collections import defaultdict
from operator import itemgetter
from typing import List, Dict, Tuple
from typing import Sequence

import numpy as np
import torch
from bert_score import BERTScorer
from nltk import PorterStemmer
from spacy.tokens import Doc, Span
from toolz import itertoolz
from transformers import AutoTokenizer
from transformers.tokenization_utils_base import PaddingStrategy


class EmbeddingModel(ABC):
    @abstractmethod
    def embed(
        self,
        sents: List[Span]
    ):
        pass


class ContextualEmbedding(EmbeddingModel):

    def __init__(self, model, tokenizer_name, max_length):
        self.model = model
        self.tokenizer = SpacyHuggingfaceTokenizer(tokenizer_name, max_length)
        self._device = model.device

    def embed(
        self,
        sents: List[Span]
    ):
        encoded_input, special_tokens_masks, token_alignments = self.tokenizer.batch_encode(sents)
        encoded_input = {k: v.to(self._device) for k, v in encoded_input.items()}
        with torch.no_grad():
            model_output = self.model(**encoded_input)
            embeddings = model_output[0].cpu()

        spacy_embs_list = []
        for embs, mask, token_alignment \
            in zip(embeddings, special_tokens_masks, token_alignments):
            mask = torch.tensor(mask)
            embs = embs[mask == 0]  # Filter embeddings at special token positions
            spacy_embs = []
            for hf_idxs in token_alignment:
                if hf_idxs is None:
                    pooled_embs = torch.zeros_like(embs[0])
                else:
                    pooled_embs = embs[hf_idxs].mean(dim=0)  # Pool embeddings that map to the same spacy token
                spacy_embs.append(pooled_embs.numpy())
            spacy_embs = np.stack(spacy_embs)
            spacy_embs = spacy_embs / np.linalg.norm(spacy_embs, axis=-1, keepdims=True)  # Normalize
            spacy_embs_list.append(spacy_embs)
        for embs, sent in zip(spacy_embs_list, sents):
            assert len(embs) == len(sent)
        return spacy_embs_list


class StaticEmbedding(EmbeddingModel):

    def embed(
        self,
        sents: List[Span]
    ):
        return [
            np.stack([t.vector / (t.vector_norm or 1) for t in sent])
            for sent in sents
        ]


class EmbeddingAligner():

    def __init__(
        self,
        embedding: EmbeddingModel,
        threshold: float,
        top_k: int,
        baseline_val=0
    ):
        self.threshold = threshold
        self.top_k = top_k
        self.embedding = embedding
        self.baseline_val = baseline_val

    def align(
        self,
        source: Doc,
        targets: Sequence[Doc]
    ) -> List[Dict]:
        """Compute alignment from summary tokens to doc tokens with greatest semantic similarity
        Args:
            source: Source spaCy document
            targets: Target spaCy documents
        Returns: List of alignments, one for each target document
        """
        if len(source) == 0:
            return [{} for _ in targets]
        all_sents = list(source.sents) + list(itertools.chain.from_iterable(target.sents for target in targets))
        chunk_sizes = [_iter_len(source.sents)] + \
                      [_iter_len(target.sents) for target in targets]
        all_sents_token_embeddings = self.embedding.embed(all_sents)
        chunked_sents_token_embeddings = _split(all_sents_token_embeddings, chunk_sizes)
        source_sent_token_embeddings = chunked_sents_token_embeddings[0]
        source_token_embeddings = np.concatenate(source_sent_token_embeddings)
        for token_idx, token in enumerate(source):
            if token.is_stop or token.is_punct:
                source_token_embeddings[token_idx] = 0
        alignments = []
        for i, target in enumerate(targets):
            target_sent_token_embeddings = chunked_sents_token_embeddings[i + 1]
            target_token_embeddings = np.concatenate(target_sent_token_embeddings)
            for token_idx, token in enumerate(target):
                if token.is_stop or token.is_punct:
                    target_token_embeddings[token_idx] = 0
            alignment = defaultdict(list)
            for score, target_idx, source_idx in self._emb_sim_sparse(
                target_token_embeddings,
                source_token_embeddings,
            ):
                alignment[target_idx].append((source_idx, score))
            # TODO used argpartition to get nlargest
            for j in list(alignment):
                alignment[j] = heapq.nlargest(self.top_k, alignment[j], itemgetter(1))
            alignments.append(alignment)
        return alignments

    def _emb_sim_sparse(self, embs_1, embs_2):
        sim = embs_1 @ embs_2.T
        sim = (sim - self.baseline_val) / (1 - self.baseline_val)
        keep = sim > self.threshold
        keep_idxs_1, keep_idxs_2 = np.where(keep)
        keep_scores = sim[keep]
        return list(zip(keep_scores, keep_idxs_1, keep_idxs_2))


class BertscoreAligner(EmbeddingAligner):
    def __init__(
        self,
        threshold,
        top_k
    ):
        scorer = BERTScorer(lang="en", rescale_with_baseline=True)
        model = scorer._model
        embedding = ContextualEmbedding(model, "roberta-large", 510)
        baseline_val = scorer.baseline_vals[2].item()

        super(BertscoreAligner, self).__init__(
            embedding, threshold, top_k, baseline_val
        )


class StaticEmbeddingAligner(EmbeddingAligner):
    def __init__(
        self,
        threshold,
        top_k
    ):
        embedding = StaticEmbedding()
        super(StaticEmbeddingAligner, self).__init__(
            embedding, threshold, top_k
        )


class NGramAligner():

    def __init__(self):
        self.stemmer = PorterStemmer()

    def align(
        self,
        source: Doc,
        targets: List[Doc],
    ) -> List[Dict]:

        alignments = []
        source_ngram_spans = self._get_ngram_spans(source)
        for target in targets:
            target_ngram_spans = self._get_ngram_spans(target)
            alignments.append(
                self._align_ngrams(target_ngram_spans, source_ngram_spans)
            )
        return alignments

    def _get_ngram_spans(
        self,
        doc: Doc,
    ):
        ngrams = []
        for sent in doc.sents:
            for n in range(1, len(list(sent))):
                tokens = [t for t in sent if not (t.is_stop or t.is_punct)]
                ngrams.extend(_ngrams(tokens, n))

        def ngram_key(ngram):
            return tuple(self.stemmer.stem(token.text).lower() for token in ngram)

        key_to_ngrams = itertoolz.groupby(ngram_key, ngrams)
        key_to_spans = {}
        for k, grouped_ngrams in key_to_ngrams.items():
            key_to_spans[k] = [
                (ngram[0].i, ngram[-1].i + 1)
                for ngram in grouped_ngrams
            ]
        return key_to_spans

    def _align_ngrams(
        self,
        ngram_spans_1: Dict[Tuple[str], List[Tuple[int, int]]],
        ngram_spans_2: Dict[Tuple[str], List[Tuple[int, int]]]
    ) -> Dict[Tuple[int, int], List[Tuple[int, int]]]:
        """Align ngram spans between two documents
        Args:
            ngram_spans_1: Map from (normalized_token1, normalized_token2, ...) n-gram tuple to a list of token spans
                of format (start_pos, end_pos)
            ngram_spans_2: Same format as above, but for second text
        Returns: map from each (start, end) span in text 1 to list of aligned (start, end) spans in text 2
        """
        if not ngram_spans_1 or not ngram_spans_2:
            return {}
        max_span_end_1 = max(span[1] for span in itertools.chain.from_iterable(ngram_spans_1.values()))
        token_is_available_1 = [True] * max_span_end_1  #
        matched_keys = list(set(ngram_spans_1.keys()) & set(ngram_spans_2.keys()))  # Matched normalized ngrams betwee
        matched_keys.sort(key=len, reverse=True)  # Process n-grams from longest to shortest

        alignment = defaultdict(list)  # Map from each matched span in text 1 to list of aligned spans in text 2
        for key in matched_keys:
            spans_1 = ngram_spans_1[key]
            spans_2 = ngram_spans_2[key]
            available_spans_1 = [span for span in spans_1 if all(token_is_available_1[slice(*span)])]
            matched_spans_1 = []
            if available_spans_1 and spans_2:
                # if ngram can be matched to available spans in both sequences
                for span in available_spans_1:
                    # It's possible that these newly matched spans may be overlapping with one another, so
                    # check that token positions still available (only one span allowed ber token in text 1):
                    if all(token_is_available_1[slice(*span)]):
                        matched_spans_1.append(span)
                        token_is_available_1[slice(*span)] = [False] * (span[1] - span[0])
            for span1 in matched_spans_1:
                alignment[span1] = spans_2

        return alignment


class SpacyHuggingfaceTokenizer:
    def __init__(
        self,
        model_name,
        max_length
    ):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
        self.max_length = max_length

    def batch_encode(
        self,
        sents: List[Span]
    ):
        token_alignments = []
        token_ids_list = []

        # Tokenize each sentence and special tokens.
        for sent in sents:
            hf_tokens, token_alignment = self.tokenize(sent)
            token_alignments.append(token_alignment)
            token_ids = self.tokenizer.convert_tokens_to_ids(hf_tokens)
            encoding = self.tokenizer.prepare_for_model(
                token_ids,
                add_special_tokens=True,
                padding=False,
            )
            token_ids_list.append(encoding['input_ids'])

        # Add padding
        max_length = max(map(len, token_ids_list))
        attention_mask = []
        input_ids = []
        special_tokens_masks = []
        for token_ids in token_ids_list:
            encoding = self.tokenizer.prepare_for_model(
                token_ids,
                padding=PaddingStrategy.MAX_LENGTH,
                max_length=max_length,
                add_special_tokens=False
            )
            input_ids.append(encoding['input_ids'])
            attention_mask.append(encoding['attention_mask'])
            special_tokens_masks.append(
                self.tokenizer.get_special_tokens_mask(
                    encoding['input_ids'],
                    already_has_special_tokens=True
                )
            )

        encoded = {
            'input_ids': torch.tensor(input_ids),
            'attention_mask': torch.tensor(attention_mask)
        }
        return encoded, special_tokens_masks, token_alignments

    def tokenize(
        self,
        sent
    ):
        """Convert spacy sentence to huggingface tokens and compute the alignment"""
        hf_tokens = []
        token_alignment = []
        for i, token in enumerate(sent):
            # "Tokenize" each word individually, so as to track the alignment between spaCy/HF tokens
            # Prefix all tokens with a space except the first one in the sentence
            if i == 0:
                token_text = token.text
            else:
                token_text = ' ' + token.text
            start_hf_idx = len(hf_tokens)
            word_tokens = self.tokenizer.tokenize(token_text)
            end_hf_idx = len(hf_tokens) + len(word_tokens)
            if end_hf_idx < self.max_length:
                hf_tokens.extend(word_tokens)
                hf_idxs = list(range(start_hf_idx, end_hf_idx))
            else:
                hf_idxs = None
            token_alignment.append(hf_idxs)
        return hf_tokens, token_alignment


def _split(data, sizes):
    it = iter(data)
    return [[next(it) for _ in range(size)] for size in sizes]


def _iter_len(it):
    return sum(1 for _ in it)

    # TODO set up batching
    # To get top K axis and value per row: https://stackoverflow.com/questions/42832711/using-np-argpartition-to-index-values-in-a-multidimensional-array


def _ngrams(tokens, n):
    for i in range(len(tokens) - n + 1):
        yield tokens[i:i + n]