wehub-resource-sync
278 lines
10 KiB
Python
278 lines
10 KiB
Python
import numpy as np
|
|
import scipy.special
|
|
|
|
from .._serializable import Deserializer, Serializer
|
|
from ..utils import safe_isinstance
|
|
from ..utils.transformers import getattr_silent
|
|
from ._model import Model
|
|
|
|
|
|
class TopKLM(Model):
|
|
"""Generates scores (log odds) for the top-k tokens for Causal/Masked LM."""
|
|
|
|
def __init__(self, model, tokenizer, k=10, generate_topk_token_ids=None, batch_size=128, device=None):
|
|
"""Take Causal/Masked LM model and tokenizer and build a log odds output model for the top-k tokens.
|
|
|
|
Parameters
|
|
----------
|
|
model: object or function
|
|
A object of any pretrained transformer model which is to be explained.
|
|
|
|
tokenizer: object
|
|
A tokenizer object(PreTrainedTokenizer/PreTrainedTokenizerFast).
|
|
|
|
generation_function_for_topk_token_ids: function
|
|
A function which is used to generate top-k token ids. Log odds will be generated for these custom token ids.
|
|
|
|
batch_size: int
|
|
Batch size for model inferencing and computing logodds (default=128).
|
|
|
|
device: str
|
|
By default, it infers if system has a gpu and accordingly sets device. Should be 'cpu' or 'cuda' or pytorch models.
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray
|
|
The scores (log odds) of generating top-k token ids using the model.
|
|
|
|
"""
|
|
super().__init__(model)
|
|
|
|
self.tokenizer = tokenizer
|
|
# set pad token if not defined
|
|
if getattr_silent(self.tokenizer, "pad_token") is None:
|
|
self.tokenizer.pad_token = self.tokenizer.eos_token
|
|
self.k = k
|
|
self._custom_generate_topk_token_ids = generate_topk_token_ids
|
|
self.batch_size = batch_size
|
|
self.device = device
|
|
|
|
self.X = None
|
|
self.topk_token_ids = None
|
|
self.output_names = None
|
|
|
|
self.model_type = None
|
|
if safe_isinstance(self.inner_model, "transformers.PreTrainedModel"):
|
|
self.model_type = "pt"
|
|
import torch
|
|
|
|
self.device = (
|
|
torch.device("cuda" if torch.cuda.is_available() else "cpu") if self.device is None else self.device
|
|
)
|
|
self.inner_model = self.inner_model.to(self.device)
|
|
elif safe_isinstance(self.inner_model, "transformers.TFPreTrainedModel"):
|
|
self.model_type = "tf"
|
|
|
|
def __call__(self, masked_X, X):
|
|
"""Computes log odds scores for a given batch of masked inputs for the top-k tokens for Causal/Masked LM.
|
|
|
|
Parameters
|
|
----------
|
|
masked_X: numpy.ndarray
|
|
An array containing a list of masked inputs.
|
|
|
|
X: numpy.ndarray
|
|
An array containing a list of original inputs
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray
|
|
A numpy array of log odds scores for top-k tokens for every input pair (masked_X, X)
|
|
|
|
"""
|
|
output_batch = None
|
|
self.update_cache_X(X[:1])
|
|
start_batch_idx, end_batch_idx = 0, len(masked_X)
|
|
while start_batch_idx < end_batch_idx:
|
|
logits = self.get_lm_logits(masked_X[start_batch_idx : start_batch_idx + self.batch_size])
|
|
logodds = self.get_logodds(logits)
|
|
if output_batch is None:
|
|
output_batch = logodds
|
|
else:
|
|
output_batch = np.concatenate((output_batch, logodds))
|
|
start_batch_idx += self.batch_size
|
|
return output_batch
|
|
|
|
def update_cache_X(self, X):
|
|
"""The function updates original input(X) and top-k token ids for the Causal/Masked LM.
|
|
|
|
It mimics the caching mechanism to update the original input and topk token ids
|
|
that are to be explained and which updates for every new row of explanation.
|
|
|
|
Parameters
|
|
----------
|
|
X: np.ndarray
|
|
Input(Text) for an explanation row.
|
|
|
|
"""
|
|
# check if the source sentence has been updated (occurs when explaining a new row)
|
|
if (self.X is None) or (not np.array_equal(self.X, X)):
|
|
self.X = X
|
|
self.output_names = self.get_output_names_and_update_topk_token_ids(self.X)
|
|
|
|
def get_output_names_and_update_topk_token_ids(self, X):
|
|
"""Gets the token names for top-k token ids for Causal/Masked LM.
|
|
|
|
Parameters
|
|
----------
|
|
X: np.ndarray
|
|
Input(Text) for an explanation row.
|
|
|
|
Returns
|
|
-------
|
|
list
|
|
A list of output tokens.
|
|
|
|
"""
|
|
# see if the user gave a custom token generator
|
|
if self._custom_generate_topk_token_ids is not None:
|
|
return self._custom_generate_topk_token_ids(X)
|
|
|
|
# otherwise we pick the top k tokens from the model
|
|
self.topk_token_ids = self.generate_topk_token_ids(X)
|
|
output_names = [self.tokenizer.decode([x]) for x in self.topk_token_ids]
|
|
return output_names
|
|
|
|
def get_logodds(self, logits):
|
|
"""Calculates log odds from logits.
|
|
|
|
This function passes the logits through softmax and then computes log odds for the top-k token ids.
|
|
|
|
Parameters
|
|
----------
|
|
logits: numpy.ndarray
|
|
An array of logits generated from the model.
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray
|
|
Computes log odds for corresponding top-k token ids.
|
|
|
|
"""
|
|
assert self.topk_token_ids is not None
|
|
|
|
# pass logits through softmax, get the token corresponding score and convert back to log odds (as one vs all)
|
|
def calc_logodds(arr):
|
|
probs = np.exp(arr) / np.exp(arr).sum(-1)
|
|
logodds = scipy.special.logit(probs)
|
|
return logodds
|
|
|
|
# pass logits through softmax, get the token corresponding score and convert back to log odds (as one vs all)
|
|
logodds = np.apply_along_axis(calc_logodds, -1, logits)
|
|
logodds_for_topk_token_ids = np.take(logodds, self.topk_token_ids, axis=-1)
|
|
return logodds_for_topk_token_ids
|
|
|
|
def get_inputs(self, X, padding_side="right"):
|
|
"""The function tokenizes source sentence.
|
|
|
|
Parameters
|
|
----------
|
|
X: numpy.ndarray
|
|
X is a batch of text.
|
|
|
|
Returns
|
|
-------
|
|
dict
|
|
Dictionary of padded source sentence ids and attention mask as tensors("pt" or "tf" based on similarity_model_type).
|
|
|
|
"""
|
|
self.tokenizer.padding_side = padding_side
|
|
inputs = self.tokenizer(X.tolist(), return_tensors=self.model_type, padding=True)
|
|
# set tokenizer padding to default
|
|
self.tokenizer.padding_side = "right"
|
|
return inputs
|
|
|
|
def generate_topk_token_ids(self, X) -> np.ndarray:
|
|
"""Generates top-k token ids for Causal/Masked LM.
|
|
|
|
Parameters
|
|
----------
|
|
X: numpy.ndarray
|
|
X is the original input sentence for an explanation row.
|
|
|
|
Returns
|
|
-------
|
|
np.ndarray
|
|
An array of top-k token ids.
|
|
|
|
"""
|
|
logits = self.get_lm_logits(X)
|
|
topk_tokens_ids = (-logits).argsort()[0, : self.k]
|
|
return topk_tokens_ids
|
|
|
|
def get_lm_logits(self, X):
|
|
"""Evaluates a Causal/Masked LM model and returns logits corresponding to next word/masked word.
|
|
|
|
Parameters
|
|
----------
|
|
X: numpy.ndarray
|
|
An array containing a list of masked inputs.
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray
|
|
Logits corresponding to next word/masked word.
|
|
|
|
"""
|
|
if self.model_type not in ["pt", "tf"]:
|
|
raise NotImplementedError("Only PyTorch and TensorFlow models are supported!")
|
|
|
|
from transformers import MODEL_FOR_CAUSAL_LM_MAPPING
|
|
|
|
if type(self.inner_model) in MODEL_FOR_CAUSAL_LM_MAPPING.values():
|
|
inputs = self.get_inputs(X, padding_side="left")
|
|
if self.model_type == "pt":
|
|
import torch
|
|
|
|
inputs["position_ids"] = inputs["attention_mask"].long().cumsum(-1) - 1
|
|
inputs["position_ids"].masked_fill_(inputs["attention_mask"] == 0, 0)
|
|
inputs = inputs.to(self.device)
|
|
# generate outputs and logits
|
|
with torch.no_grad():
|
|
outputs = self.inner_model(**inputs, return_dict=True)
|
|
# extract only logits corresponding to target sentence ids
|
|
logits = outputs.logits.detach().cpu().numpy().astype("float64")[:, -1, :]
|
|
else:
|
|
assert self.model_type == "tf"
|
|
import tensorflow as tf
|
|
|
|
inputs["position_ids"] = tf.math.cumsum(inputs["attention_mask"], axis=-1) - 1
|
|
inputs["position_ids"] = tf.where(inputs["attention_mask"] == 0, 0, inputs["position_ids"])
|
|
if self.device is None:
|
|
outputs = self.inner_model(inputs, return_dict=True)
|
|
else:
|
|
try:
|
|
with tf.device(self.device):
|
|
outputs = self.inner_model(inputs, return_dict=True)
|
|
except RuntimeError as err:
|
|
print(err)
|
|
logits = outputs.logits.numpy().astype("float64")[:, -1, :]
|
|
else:
|
|
raise NotImplementedError(f"Model type '{type(self.inner_model)}' not supported!")
|
|
return logits
|
|
|
|
def save(self, out_file):
|
|
super().save(out_file)
|
|
|
|
# Increment the version number when the encoding changes!
|
|
with Serializer(out_file, "shap.models.TextGeneration", version=0) as s:
|
|
s.save("tokenizer", self.tokenizer)
|
|
s.save("k", self.k)
|
|
s.save("generate_topk_token_ids", self._custom_generate_topk_token_ids)
|
|
s.save("batch_size", self.batch_size)
|
|
s.save("device", self.device)
|
|
|
|
@classmethod
|
|
def load(cls, in_file, instantiate=True):
|
|
if instantiate:
|
|
return cls._instantiated_load(in_file)
|
|
|
|
kwargs = super().load(in_file, instantiate=False)
|
|
with Deserializer(in_file, "shap.models.TextGeneration", min_version=0, max_version=0) as s:
|
|
kwargs["tokenizer"] = s.load("tokenizer")
|
|
kwargs["k"] = s.load("k")
|
|
kwargs["generate_topk_token_ids"] = s.load("generate_topk_token_ids")
|
|
kwargs["batch_size"] = s.load("batch_size")
|
|
kwargs["device"] = s.load("device")
|
|
return kwargs
|