wehub-resource-sync
443 lines
19 KiB
Python
443 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import inspect
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import scipy.special
|
|
|
|
from .. import models
|
|
from .._serializable import Deserializer, Serializer
|
|
from ..utils import safe_isinstance
|
|
from ..utils.transformers import getattr_silent, parse_prefix_suffix_for_tokenizer
|
|
from ._model import Model
|
|
|
|
|
|
class TeacherForcing(Model):
|
|
"""Generates scores (log odds) for output text explanation algorithms using Teacher Forcing technique.
|
|
|
|
This class supports generation of log odds for transformer models as well as functions. In model agnostic
|
|
cases (model is function) it expects a similarity_model and similarity_tokenizer to approximate log odd scores
|
|
for target sentence generated by the model.
|
|
"""
|
|
|
|
def __init__(
|
|
self, model, tokenizer=None, similarity_model=None, similarity_tokenizer=None, batch_size=128, device=None
|
|
):
|
|
"""Build a teacher forcing model from the given text generation model.
|
|
|
|
Parameters
|
|
----------
|
|
model: object or function
|
|
A object of any pretrained transformer model or function which is to be explained.
|
|
|
|
tokenizer: object
|
|
A tokenizer object(PreTrainedTokenizer/PreTrainedTokenizerFast) which is used to tokenize source and target sentence.
|
|
|
|
similarity_model: object
|
|
A pretrained transformer model object which is used in model agnostic scenario to approximate log odds.
|
|
|
|
similarity_tokenizer: object
|
|
A tokenizer object(PreTrainedTokenizer/PreTrainedTokenizerFast) which is used to tokenize sentence in model agnostic scenario.
|
|
|
|
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 target sentence ids using the model.
|
|
|
|
"""
|
|
super().__init__(model)
|
|
|
|
self.tokenizer = tokenizer
|
|
# set pad token if not defined
|
|
if self.tokenizer is not None and getattr_silent(self.tokenizer, "pad_token") is None:
|
|
self.tokenizer.pad_token = self.tokenizer.eos_token
|
|
|
|
# set our working device
|
|
self.device = device
|
|
if self.device is None:
|
|
if getattr(model, "device", None) is not None:
|
|
self.device = model.device
|
|
elif getattr(similarity_model, "device", None) is not None:
|
|
self.device = similarity_model.device
|
|
|
|
self.batch_size = batch_size
|
|
# assign text generation function
|
|
if safe_isinstance(model, "transformers.PreTrainedModel") or safe_isinstance(
|
|
model, "transformers.TFPreTrainedModel"
|
|
):
|
|
self.text_generate = models.TextGeneration(self.inner_model, tokenizer=self.tokenizer, device=self.device)
|
|
self.similarity_model = model
|
|
self.similarity_tokenizer = tokenizer
|
|
self.model_agnostic = False
|
|
else:
|
|
self.text_generate = models.TextGeneration(self.inner_model, device=self.device)
|
|
self.similarity_model = similarity_model
|
|
self.similarity_tokenizer = similarity_tokenizer
|
|
# set pad token for a similarity tokenizer(in a model agnostic scenario) if not defined
|
|
if self.similarity_tokenizer is not None and self.similarity_tokenizer.pad_token is None:
|
|
self.similarity_tokenizer.pad_token = self.similarity_tokenizer.eos_token
|
|
self.model_agnostic = True
|
|
# initializing target which is the target sentence/ids for every new row of explanation
|
|
self.output: np.ndarray | None = None
|
|
self.output_names: list[Any] | None = None
|
|
|
|
self.similarity_model_type = None
|
|
if safe_isinstance(self.similarity_model, "transformers.PreTrainedModel"):
|
|
self.similarity_model_type = "pt"
|
|
if (
|
|
self.device is not None
|
|
): # = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if self.device is None else self.device
|
|
d = self.similarity_model.device
|
|
assert d == self.device or str(d) == self.device, (
|
|
"The passed similarity_model must be on the same device!"
|
|
)
|
|
# self.similarity_model = self.similarity_model.to(self.device)
|
|
elif safe_isinstance(self.similarity_model, "transformers.TFPreTrainedModel"):
|
|
self.similarity_model_type = "tf"
|
|
|
|
def __call__(self, X, Y):
|
|
"""Computes log odds scores of generating output(text) for a given batch of input(text/image) .
|
|
|
|
Parameters
|
|
----------
|
|
X: numpy.ndarray
|
|
An array containing a list of masked inputs.
|
|
|
|
Y: numpy.ndarray
|
|
An array containing a list of target sentence/ids.
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray
|
|
A numpy array of log odds scores for every input pair (masked_X, X)
|
|
|
|
"""
|
|
output_batch = None
|
|
# caching updates output names and target sentence ids
|
|
self.update_output_names(Y[:1])
|
|
start_batch_idx, end_batch_idx = 0, len(X)
|
|
while start_batch_idx < end_batch_idx:
|
|
X_batch = X[start_batch_idx : start_batch_idx + self.batch_size]
|
|
Y_batch = Y[start_batch_idx : start_batch_idx + self.batch_size]
|
|
logits = self.get_teacher_forced_logits(X_batch, Y_batch)
|
|
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_output_names(self, output: np.ndarray):
|
|
"""The function updates output tokens.
|
|
|
|
It mimics the caching mechanism to update the output tokens for every
|
|
new row of explanation that are to be explained.
|
|
|
|
Parameters
|
|
----------
|
|
output: numpy.ndarray
|
|
Output(sentence/sentence ids) for an explanation row.
|
|
|
|
"""
|
|
# check if the target sentence has been updated (occurs when explaining a new row)
|
|
if (self.output is None) or (not np.array_equal(self.output, output)):
|
|
self.output = output
|
|
self.output_names = self.get_output_names(output)
|
|
|
|
def get_output_names(self, output):
|
|
"""Gets the output tokens by computing the output sentence ids and output names using the similarity_tokenizer.
|
|
|
|
Parameters
|
|
----------
|
|
output: numpy.ndarray
|
|
Output(sentence/sentence ids) for an explanation row.
|
|
|
|
Returns
|
|
-------
|
|
list
|
|
A list of output tokens.
|
|
|
|
"""
|
|
output_ids = self.get_outputs(output)
|
|
output_names = [self.similarity_tokenizer.decode([x]).strip() for x in output_ids[0, :]]
|
|
return output_names
|
|
|
|
def get_outputs(self, X):
|
|
"""The function tokenizes output sentences and returns ids.
|
|
|
|
Parameters
|
|
----------
|
|
X: numpy.ndarray
|
|
Output(sentence/sentence ids) for an explanation row.
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray
|
|
An array of output(target sentence) ids.
|
|
|
|
"""
|
|
# check if output is a sentence or already parsed target ids
|
|
if X.dtype.type is np.str_:
|
|
parsed_tokenizer_dict = parse_prefix_suffix_for_tokenizer(self.similarity_tokenizer)
|
|
keep_prefix, keep_suffix = parsed_tokenizer_dict["keep_prefix"], parsed_tokenizer_dict["keep_suffix"]
|
|
if keep_suffix > 0:
|
|
output_ids = np.array(self.similarity_tokenizer(X.tolist(), padding=True)["input_ids"])[
|
|
:, keep_prefix:-keep_suffix
|
|
]
|
|
else:
|
|
output_ids = np.array(self.similarity_tokenizer(X.tolist(), padding=True)["input_ids"])[:, keep_prefix:]
|
|
else:
|
|
output_ids = X
|
|
return output_ids
|
|
|
|
def get_inputs(self, X, padding_side="right"):
|
|
"""The function tokenizes source sentences.
|
|
|
|
In model agnostic case, the function calls model(X) which is expected to
|
|
return a batch of output sentences which is tokenized to compute inputs.
|
|
|
|
Parameters
|
|
----------
|
|
X: numpy.ndarray
|
|
X could be a batch of text or images(model agnostic case).
|
|
|
|
Returns
|
|
-------
|
|
dict
|
|
Dictionary of padded source sentence ids and attention mask as tensors("pt" or "tf" based on similarity_model_type).
|
|
|
|
"""
|
|
if self.model_agnostic:
|
|
# In model agnostic case, we first pass the input through the model and then tokenize output sentence
|
|
input_sentences = np.array(self.inner_model(X))
|
|
else:
|
|
input_sentences = np.array(X)
|
|
# set tokenizer padding to prepare inputs for batch inferencing
|
|
# padding_side="left" for only decoder models text generation eg. GPT2
|
|
self.similarity_tokenizer.padding_side = padding_side
|
|
inputs = self.similarity_tokenizer(
|
|
input_sentences.tolist(), return_tensors=self.similarity_model_type, padding=True
|
|
)
|
|
# set tokenizer padding to default
|
|
self.similarity_tokenizer.padding_side = "right"
|
|
return inputs
|
|
|
|
def get_logodds(self, logits):
|
|
"""Calculates log odds from logits.
|
|
|
|
This function passes the logits through softmax and then computes log odds for the output(target sentence) ids.
|
|
|
|
Parameters
|
|
----------
|
|
logits: numpy.ndarray
|
|
An array of logits generated from the model.
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray
|
|
Computes log odds for corresponding output ids.
|
|
|
|
"""
|
|
# set output ids for which scores are to be extracted
|
|
assert self.output is not None
|
|
if self.output.dtype.type is np.str_:
|
|
output_ids = self.get_outputs(self.output)[0]
|
|
else:
|
|
output_ids = self.output[0]
|
|
|
|
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_output_ids = logodds[:, np.array(range(logodds.shape[1])), output_ids]
|
|
return logodds_for_output_ids
|
|
|
|
def model_inference(self, inputs, output_ids):
|
|
"""This function performs model inference for tensorflow and pytorch models.
|
|
|
|
Parameters
|
|
----------
|
|
inputs: dict
|
|
Dictionary of padded source sentence ids and attention mask as tensors.
|
|
|
|
output_ids: numpy.ndarray
|
|
An array of decoder output ids.
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray
|
|
Returns output logits from the model.
|
|
|
|
"""
|
|
if self.similarity_model_type == "pt":
|
|
import torch
|
|
|
|
# create torch tensors and move to device
|
|
if self.device is not None:
|
|
inputs = inputs.to(self.device)
|
|
output_ids = torch.tensor(output_ids, dtype=torch.int64, device=self.device)
|
|
self.similarity_model.eval()
|
|
with torch.no_grad():
|
|
if self.similarity_model.config.is_encoder_decoder:
|
|
# model inference
|
|
outputs = self.similarity_model(
|
|
**inputs, decoder_input_ids=output_ids, labels=output_ids, return_dict=True
|
|
)
|
|
else:
|
|
# combine source and target sentence ids to pass into decoder eg: in case of distillgpt2
|
|
inputs["input_ids"] = torch.cat((inputs["input_ids"], output_ids), dim=-1)
|
|
attention_mask_for_output_ids = torch.ones(
|
|
output_ids.shape, dtype=output_ids.dtype, device=self.device
|
|
)
|
|
inputs["attention_mask"] = torch.cat(
|
|
(inputs["attention_mask"], attention_mask_for_output_ids), dim=-1
|
|
)
|
|
# create position ids due to left padding for decoder models
|
|
inputs["position_ids"] = inputs["attention_mask"].long().cumsum(-1) - 1
|
|
inputs["position_ids"].masked_fill_(inputs["attention_mask"] == 0, 0)
|
|
# model inference
|
|
expected_parameters = list(inspect.signature(self.similarity_model.forward).parameters)
|
|
inputs = {k: v for k, v in inputs.items() if k in expected_parameters}
|
|
outputs = self.similarity_model(**inputs, return_dict=True)
|
|
logits = outputs.logits.detach().cpu().numpy().astype("float64")
|
|
elif self.similarity_model_type == "tf":
|
|
import tensorflow as tf
|
|
|
|
output_ids = tf.convert_to_tensor(output_ids, dtype=tf.int32)
|
|
if self.similarity_model.config.is_encoder_decoder:
|
|
if self.device is None:
|
|
outputs = self.similarity_model(
|
|
inputs, decoder_input_ids=output_ids, labels=output_ids, return_dict=True
|
|
)
|
|
else:
|
|
try:
|
|
with tf.device(self.device):
|
|
outputs = self.similarity_model(
|
|
inputs, decoder_input_ids=output_ids, labels=output_ids, return_dict=True
|
|
)
|
|
except RuntimeError as err:
|
|
print(err)
|
|
else:
|
|
# combine source and target sentence ids to pass into decoder eg: in case of distillgpt2
|
|
inputs["input_ids"] = tf.concat((inputs["input_ids"], output_ids), axis=-1)
|
|
attention_mask_for_output_ids = tf.ones(output_ids.shape, dtype=output_ids.dtype)
|
|
inputs["attention_mask"] = tf.concat((inputs["attention_mask"], attention_mask_for_output_ids), axis=-1)
|
|
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.similarity_model(inputs, return_dict=True)
|
|
else:
|
|
try:
|
|
with tf.device(self.device):
|
|
outputs = self.similarity_model(inputs, return_dict=True)
|
|
except RuntimeError as err:
|
|
print(err)
|
|
logits = outputs.logits.numpy().astype("float64")
|
|
return logits
|
|
|
|
def get_teacher_forced_logits(self, X, Y):
|
|
"""The function generates logits for transformer models.
|
|
|
|
It generates logits for encoder-decoder models as well as decoder only models by using the teacher forcing technique.
|
|
|
|
Parameters
|
|
----------
|
|
X: numpy.ndarray
|
|
An array containing a list of masked inputs.
|
|
|
|
Y: numpy.ndarray
|
|
An array containing a list of target sentence/ids.
|
|
|
|
Returns
|
|
-------
|
|
numpy.ndarray
|
|
Decoder output logits for output(target sentence) ids.
|
|
|
|
"""
|
|
# check if type of model architecture assigned in model config
|
|
if (
|
|
hasattr(self.similarity_model.config, "is_encoder_decoder")
|
|
and not self.similarity_model.config.is_encoder_decoder
|
|
) and (hasattr(self.similarity_model.config, "is_decoder") and not self.similarity_model.config.is_decoder):
|
|
pass # self.similarity_model.config.is_decoder = True # TODOmaybe: is this okay?
|
|
# raise ValueError(
|
|
# "Please assign either of is_encoder_decoder or is_decoder to True in model config for extracting target sentence ids"
|
|
# )
|
|
# get output ids for teacher forcing
|
|
output_ids = self.get_outputs(Y)
|
|
if self.similarity_model.config.is_encoder_decoder:
|
|
# encode batched inputs by padding on the right side
|
|
inputs = self.get_inputs(X, padding_side="right")
|
|
# assigning decoder start token id as it is needed for encoder decoder model generation
|
|
decoder_start_token_id = None
|
|
if (
|
|
hasattr(self.similarity_model.config, "decoder_start_token_id")
|
|
and self.similarity_model.config.decoder_start_token_id is not None
|
|
):
|
|
decoder_start_token_id = self.similarity_model.config.decoder_start_token_id
|
|
elif (
|
|
hasattr(self.similarity_model.config, "bos_token_id")
|
|
and self.similarity_model.config.bos_token_id is not None
|
|
):
|
|
decoder_start_token_id = self.similarity_model.config.bos_token_id
|
|
elif (
|
|
hasattr(self.similarity_model.config, "decoder")
|
|
and hasattr(self.similarity_model.config.decoder, "bos_token_id")
|
|
and self.similarity_model.config.decoder.bos_token_id is not None
|
|
):
|
|
decoder_start_token_id = self.similarity_model.config.decoder.bos_token_id
|
|
else:
|
|
raise ValueError(
|
|
"No decoder_start_token_id or bos_token_id defined in config for encoder-decoder generation"
|
|
)
|
|
# concat decoder start token id to target sentence ids
|
|
output_start_id = np.ones((output_ids.shape[0], 1)) * decoder_start_token_id
|
|
output_ids = np.concatenate((output_start_id, output_ids), axis=-1)
|
|
# generate outputs and logits
|
|
logits = self.model_inference(inputs, output_ids)
|
|
logits = logits[:, :-1, :]
|
|
else:
|
|
# encode batched inputs by padding on the left side
|
|
inputs = self.get_inputs(X, padding_side="left")
|
|
# generate outputs and logits
|
|
logits = self.model_inference(inputs, output_ids)
|
|
# extract only logits corresponding to target sentence ids
|
|
logits = logits[:, -output_ids.shape[1] - 1 : -1, :]
|
|
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.TeacherForcing", version=0) as s:
|
|
s.save("tokenizer", self.tokenizer)
|
|
s.save("similarity_model", self.similarity_model)
|
|
s.save("similarity_tokenizer", self.similarity_tokenizer)
|
|
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.TeacherForcing", min_version=0, max_version=0) as s:
|
|
kwargs["tokenizer"] = s.load("tokenizer")
|
|
kwargs["similarity_model"] = s.load("similarity_model")
|
|
kwargs["similarity_tokenizer"] = s.load("similarity_tokenizer")
|
|
kwargs["batch_size"] = s.load("batch_size")
|
|
kwargs["device"] = s.load("device")
|
|
return kwargs
|