align.py

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
from abc import ABC

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, batch_size=32):
        self.model = model
        self.tokenizer = SpacyHuggingfaceTokenizer(tokenizer_name, max_length)
        self._device = model.device
        self.batch_size = batch_size

    def embed(
        self,
        sents: List[Span]
    ):
        spacy_embs_list = []
        for start_idx in range(0, len(sents), self.batch_size):
            batch = sents[start_idx: start_idx + self.batch_size]
            encoded_input, special_tokens_masks, token_alignments = self.tokenizer.batch_encode(batch)
            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()
            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 Aligner(ABC):
    @abstractmethod
    def align(
        self,
        source: Doc,
        targets: Sequence[Doc]
    ) -> List[Dict]:
        """Compute alignment from summary tokens to doc tokens
        Args:
            source: Source spaCy document
            targets: Target spaCy documents
        Returns: List of alignments, one for each target document"""
        pass


class EmbeddingAligner(Aligner):

    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(Aligner):

    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]