File size: 5,532 Bytes
fda57dd |
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 |
"""
Source: https://github.com/ZurichNLP/recognizing-semantic-differences
MIT License
Copyright (c) 2023 University of Zurich
"""
import itertools
from typing import List, Union
import torch
import transformers
from transformers import FeatureExtractionPipeline, Pipeline
from recognizers.base import DifferenceRecognizer
from recognizers.utils import DifferenceSample
Ngram = List[int] # A span of subword indices
class FeatureExtractionRecognizer(DifferenceRecognizer):
def __init__(self,
model_name_or_path: str = None,
pipeline: Union[FeatureExtractionPipeline, Pipeline] = None,
layer: int = -1,
batch_size: int = 16,
):
assert model_name_or_path is not None or pipeline is not None
if pipeline is None:
pipeline = transformers.pipeline(
model=model_name_or_path,
task="feature-extraction",
)
self.pipeline = pipeline
self.layer = layer
self.batch_size = batch_size
def encode_batch(self, sentences: List[str], **kwargs) -> torch.Tensor:
model_inputs = self.pipeline.tokenizer(sentences, return_tensors="pt", padding=True, truncation=True)
model_inputs = model_inputs.to(self.pipeline.device)
outputs = self.pipeline.model(**model_inputs, output_hidden_states=True, **kwargs)
return outputs.hidden_states[self.layer]
def predict(self,
a: str,
b: str,
**kwargs,
) -> DifferenceSample:
return self.predict_all([a], [b], **kwargs)[0]
def predict_all(self,
a: List[str],
b: List[str],
**kwargs,
) -> List[DifferenceSample]:
samples = []
for i in range(0, len(a), self.batch_size):
samples.extend(self._predict_all(
a[i:i + self.batch_size],
b[i:i + self.batch_size],
**kwargs,
))
return samples
@torch.no_grad()
def _predict_all(self,
a: List[str],
b: List[str],
**kwargs,
) -> List[DifferenceSample]:
raise NotImplementedError
def _pool(self, token_embeddings: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
"""
:param token_embeddings: batch x seq_len x dim
:param mask: batch x seq_len; 1 if token should be included in the pooling
:return: batch x dim
Do only sum and do not divide by the number of tokens because cosine similarity is length-invariant.
"""
return torch.sum(token_embeddings * mask.unsqueeze(-1), dim=1)
def _get_subwords_by_word(self, sentence: str) -> List[Ngram]:
"""
:return: For each word in the sentence, the positions of the subwords that make up the word.
"""
batch_encoding = self.pipeline.tokenizer(
sentence,
padding=True,
truncation=True,
)
subword_ids: List[List[int]] = []
for subword_idx in range(len(batch_encoding.encodings[0].word_ids)):
if batch_encoding.encodings[0].word_ids[subword_idx] is None: # Special token
continue
char_idx = batch_encoding.encodings[0].offsets[subword_idx][0]
if isinstance(self.pipeline.tokenizer, transformers.XLMRobertaTokenizerFast) or \
isinstance(self.pipeline.tokenizer, transformers.XLMRobertaTokenizer):
token = batch_encoding.encodings[0].tokens[subword_idx]
is_tail = not token.startswith("▁") and token not in self.pipeline.tokenizer.all_special_tokens
elif isinstance(self.pipeline.tokenizer, transformers.RobertaTokenizerFast) or \
isinstance(self.pipeline.tokenizer, transformers.RobertaTokenizer):
token = batch_encoding.encodings[0].tokens[subword_idx]
is_tail = not token.startswith("Ġ") and token not in self.pipeline.tokenizer.all_special_tokens
else:
is_tail = char_idx > 0 and char_idx == batch_encoding.encodings[0].offsets[subword_idx - 1][1]
if is_tail and len(subword_ids) > 0:
subword_ids[-1].append(subword_idx)
else:
subword_ids.append([subword_idx])
return subword_ids
def _get_ngrams(self, subwords_by_word: List[Ngram]) -> List[Ngram]:
"""
:return: For each subword ngram in the sentence, the positions of the subwords that make up the ngram.
"""
subwords = list(itertools.chain.from_iterable(subwords_by_word))
# Always return at least one ngram (reduce n if necessary)
min_n = min(self.min_n, len(subwords))
ngrams = []
for n in range(min_n, self.max_n + 1):
for i in range(len(subwords) - n + 1):
ngrams.append(subwords[i:i + n])
return ngrams
def _subword_labels_to_word_labels(self, subword_labels: torch.Tensor, subwords_by_words: List[Ngram]) -> List[float]:
"""
:param subword_labels: num_subwords
:param subwords_by_words: num_words x num_subwords
:return: num_words
"""
labels = []
for subword_indices in subwords_by_words:
label = subword_labels[subword_indices].mean().item()
labels.append(label)
return labels
|