def run_with_metadata(
self,
indexed_inputs: Sequence[IndexedInput],
model: lit_model.Model,
dataset: lit_dataset.IndexedDataset,
model_outputs: Optional[List[JsonDict]] = None,
config: Optional[JsonDict] = None) -> Optional[List[JsonDict]]:
"""Finds the nearest neighbors of the example specified in the config.
Args:
indexed_inputs: the dataset example to find nearest neighbors for.
model: the model being explained.
dataset: the dataset which the current examples belong to.
model_outputs: optional model outputs from calling model.predict(inputs).
config: a config which should specify:
{
'num_neighbors': [the number of nearest neighbors to return]
'dataset_name': [the name of the dataset (used for caching)]
'embedding_name': [the name of the embedding field to use]
}
Returns:
A JsonDict containing the a list of num_neighbors nearest neighbors,
where each has the example id and distance from the main example.
"""
config = NearestNeighborsConfig(**config)
dataset_outputs = list(
model.predict_with_metadata(dataset.indexed_examples,
dataset_name=config.dataset_name))
example_outputs = list(
model.predict_with_metadata(indexed_inputs,
dataset_name=config.dataset_name))
# TODO(lit-dev): Add support for selecting nearest neighbors of a set.
if len(example_outputs) != 1:
raise ValueError('More than one selected example was passed in.')
example_output = example_outputs[0]
# <float32>[emb_size]
dataset_embs = [
output[config.embedding_name] for output in dataset_outputs
]
example_embs = [example_output[config.embedding_name]]
distances = distance.cdist(example_embs, dataset_embs)[0]
sorted_indices = np.argsort(distances)
k = config.num_neighbors
k_nearest_neighbors = [{
'id':
dataset.indexed_examples[original_index]['id'],
'nn_distance':
distances[original_index]
} for original_index in sorted_indices[:k]]
return [{'nearest_neighbors': k_nearest_neighbors}]
def _get_embedding(self, example: JsonDict, model: lit_model.Model,
dataset: lit_dataset.IndexedDataset,
embedding_name: str, dataset_name: str):
"""Calls the model on the example to get the embedding."""
model_input = dataset.index_inputs([example])
model_output = model.predict_with_metadata(model_input,
dataset_name=dataset_name)
embedding = list(model_output)[0][embedding_name]
return embedding
def run_with_metadata(self,
indexed_inputs: Sequence[IndexedInput],
model: lit_model.Model,
dataset: lit_dataset.IndexedDataset,
model_outputs: Optional[List[JsonDict]] = None,
config: Optional[JsonDict] = None) -> List[JsonDict]:
if model_outputs is None:
model_outputs = list(model.predict_with_metadata(indexed_inputs))
# TODO(lit-team): pre-compute this mapping in constructor?
# This would require passing a model name to this function so we can
# reference a pre-computed list.
spec = model.spec()
field_map = map_pred_keys(dataset.spec(), spec.output,
self.is_compatible)
ret = []
for pred_key, label_key in field_map.items():
# Extract fields
labels = [ex['data'][label_key] for ex in indexed_inputs]
preds = [mo[pred_key] for mo in model_outputs]
indices = [ex['id'] for ex in indexed_inputs]
metas = [ex.get('meta', {}) for ex in indexed_inputs]
# Compute metrics, as dict(str -> float)
metrics = self.compute_with_metadata(
labels,
preds,
label_spec=dataset.spec()[label_key],
pred_spec=spec.output[pred_key],
indices=indices,
metas=metas,
config=config.get(pred_key) if config else None)
# NaN is not a valid JSON value, so replace with None which will be
# serialized as null.
# TODO(lit-team): move this logic into serialize.py somewhere instead?
metrics = {
k: (v if not np.isnan(v) else None)
for k, v in metrics.items()
}
# Format for frontend.
ret.append({
'pred_key': pred_key,
'label_key': label_key,
'metrics': metrics
})
return ret
def _train_instance(self, model: lit_model.Model,
dataset: lit_dataset.IndexedDataset, config: JsonDict,
name: Text) -> ProjectionInterpreter:
# Ignore pytype warning about abstract methods, since this should always
# be a subclass of ProjectorModel which has these implemented.
projector = self._model_factory(**config.get("proj_kw", {})) # pytype: disable=not-instantiable
train_inputs = dataset.indexed_examples
# TODO(lit-dev): remove 'dataset_name' from caching logic so we don't need
# to track it here or elsewhere.
train_outputs = list(
model.predict_with_metadata(
train_inputs, dataset_name=config.get("dataset_name")))
logging.info("Creating new projection instance on %d points",
len(train_inputs))
return ProjectionInterpreter(model,
train_inputs,
train_outputs,
projector=projector,
field_name=config["field_name"],
name=name)
def _filter_ds_examples(
self,
dataset: lit_dataset.IndexedDataset,
dataset_name: Text,
model: lit_model.Model,
reference_output: JsonDict,
pred_key: Text,
regression_thresh: Optional[float] = None) -> List[JsonDict]:
"""Reads all dataset examples and returns only those that are flips."""
if not isinstance(dataset, lit_dataset.IndexedDataset):
raise ValueError(
'Only indexed datasets are currently supported by the TabularMTC'
'generator.')
indexed_examples = list(dataset.indexed_examples)
filtered_examples = []
preds = model.predict_with_metadata(indexed_examples,
dataset_name=dataset_name)
# Find all DS examples that are flips with respect to the reference example.
for indexed_example, pred in zip(indexed_examples, preds):
flip = cf_utils.is_prediction_flip(
cf_output=pred,
orig_output=reference_output,
output_spec=model.output_spec(),
pred_key=pred_key,
regression_thresh=regression_thresh)
if flip:
candidate_example = indexed_example['data'].copy()
self._find_dataset_parent_and_set(
model_output_spec=model.output_spec(),
pred_key=pred_key,
dataset_spec=dataset.spec(),
example=candidate_example,
predicted_value=pred[pred_key])
filtered_examples.append(candidate_example)
return filtered_examples
def _train_instance(self, model: lit_model.Model,
dataset: lit_dataset.Dataset, config: Dict[Text, Any],
name: Text) -> ProjectionInterpreter:
# Ignore pytype warning about abstract methods, since this should always
# be a subclass of ProjectorModel which has these implemented.
projector = self._model_factory(**config.get("proj_kw", {})) # pytype: disable=not-instantiable
# TODO(lit-dev): recomputing hashes here is a bit wasteful - consider
# creating an 'IndexedDataset' class in the server, and passing that
# around so that components can access IndexedInputs directly.
train_inputs = caching.add_hashes_to_input(dataset.examples)
# TODO(lit-dev): remove 'dataset_name' from caching logic so we don't need
# to track it here or elsewhere.
train_outputs = list(
model.predict_with_metadata(
train_inputs, dataset_name=config.get("dataset_name")))
logging.info("Creating new projection instance on %d points",
len(train_inputs))
return ProjectionInterpreter(
model,
train_inputs,
train_outputs,
projector=projector,
field_name=config["field_name"],
name=name)
def run_with_metadata(
self,
indexed_inputs: Sequence[IndexedInput],
model: lit_model.Model,
dataset: lit_dataset.IndexedDataset,
model_outputs: Optional[List[JsonDict]] = None,
config: Optional[JsonDict] = None) -> Optional[List[JsonDict]]:
"""Runs the TCAV method given the params in the inputs and config.
Args:
indexed_inputs: all examples in the dataset, in the indexed input format.
model: the model being explained.
dataset: the dataset which the current examples belong to.
model_outputs: optional model outputs from calling model.predict(inputs).
config: a config which should specify: {
'concept_set_ids': [list of ids to use in concept set]
'class_to_explain': [gradient class to explain],
'grad_layer': [the Gradient field key of the layer to explain],
'random_state': [an optional seed to make outputs deterministic]
'dataset_name': [the name of the dataset (used for caching)]
'test_size': [Percentage of the example set to use in the LM test set]
'negative_set_ids': [optional list of ids to use as negative set] }
Returns:
A JsonDict containing the TCAV scores, directional derivatives,
statistical test p-values, and LM accuracies.
"""
config = TCAVConfig(**config)
# TODO(b/171513556): get these from the Dataset object once indices are
# available there.
dataset_examples = indexed_inputs
# Get this layer's output spec keys for gradients and embeddings.
grad_layer = config.grad_layer
output_spec = model.output_spec()
emb_layer = cast(types.Gradients, output_spec[grad_layer]).grad_for
# Get the class that the gradients were computed for.
grad_class_key = cast(types.Gradients,
output_spec[grad_layer]).grad_target_field_key
ids_set = set(config.concept_set_ids)
concept_set = [ex for ex in dataset_examples if ex['id'] in ids_set]
non_concept_set = [
ex for ex in dataset_examples if ex['id'] not in ids_set
]
# Get outputs using model.predict().
dataset_outputs = list(
model.predict_with_metadata(dataset_examples,
dataset_name=config.dataset_name))
if config.negative_set_ids:
negative_ids_set = set(config.negative_set_ids)
negative_set = [
ex for ex in dataset_examples if ex['id'] in negative_ids_set
]
return self._run_relative_tcav(grad_layer, emb_layer,
grad_class_key, concept_set,
negative_set, dataset_outputs,
model, config)
else:
return self._run_default_tcav(grad_layer, emb_layer,
grad_class_key, concept_set,
non_concept_set, dataset_outputs,
model, config)
def run_with_metadata(
self,
indexed_inputs: Sequence[IndexedInput],
model: lit_model.Model,
dataset: lit_dataset.IndexedDataset,
model_outputs: Optional[List[JsonDict]] = None,
config: Optional[JsonDict] = None) -> Optional[List[JsonDict]]:
"""Runs the TCAV method given the params in the inputs and config.
Args:
indexed_inputs: all examples in the dataset, in the indexed input format.
model: the model being explained.
dataset: the dataset which the current examples belong to.
model_outputs: optional model outputs from calling model.predict(inputs).
config: a config which should specify:
{
'concept_set_ids': [list of ids to use in concept set]
'class_to_explain': [gradient class to explain],
'grad_layer': [the Gradient field key of the layer to explain],
'random_state': [an optional seed to make outputs deterministic]
}
Returns:
A JsonDict containing the TCAV scores, directional derivatives,
statistical test p-values, and LM accuracies.
"""
config = TCAVConfig(**config)
# TODO(b/171513556): get these from the Dataset object once indices are
# available there.
dataset_examples = indexed_inputs
# Get this layer's output spec keys for gradients and embeddings.
grad_layer = config.grad_layer
output_spec = model.output_spec()
emb_layer = cast(types.Gradients, output_spec[grad_layer]).grad_for
# Get the class that the gradients were computed for.
grad_class_key = cast(types.Gradients, output_spec[grad_layer]).grad_target
ids_set = set(config.concept_set_ids)
concept_set = [ex for ex in dataset_examples if ex['id'] in ids_set]
non_concept_set = [ex for ex in dataset_examples if ex['id'] not in ids_set]
# Get outputs using model.predict().
dataset_outputs = list(model.predict_with_metadata(dataset_examples))
def _subsample(examples, n):
return random.sample(examples, n) if n < len(examples) else examples
concept_outputs = list(model.predict_with_metadata(concept_set))
non_concept_outputs = list(model.predict_with_metadata(non_concept_set))
concept_results = []
# If there are more concept set examples than non-concept set examples, we
# use random splits of the concept examples as the concept set and use the
# remainder of the dataset as the comparison set. Otherwise, we use random
# splits of the non-concept examples as the comparison set.
n = min(len(concept_set), len(non_concept_set))
# If there are an equal number of concept and non-concept examples, we
# decrease n by one so that we also sample a different set in each TCAV run.
if len(concept_set) == len(non_concept_set):
n -= 1
for _ in range(NUM_SPLITS):
concept_split_outputs = _subsample(concept_outputs, n)
comparison_split_outputs = _subsample(non_concept_outputs, n)
concept_results.append(self._run_tcav(concept_split_outputs,
comparison_split_outputs,
dataset_outputs,
config.class_to_explain,
emb_layer,
grad_layer,
grad_class_key,
config.test_size,
config.random_state))
cav_scores = [res['score'] for res in concept_results]
p_val = self.hyp_test(cav_scores)
# Get index of CAV result with the highest accuracy.
accuracies = [res['accuracy'] for res in concept_results]
index = np.argmax(accuracies)
# Many CAVS are trained and checked for statistical testing to calculate
# the p-value. The values of the first CAV are returned.
results = {'result': concept_results[index], 'p_val': p_val}
return [results]
