<jupyter_start><jupyter_text><jupyter_code>% pip install shap
import tensorflow as tf
from tensorflow.keras import layers
import tensorflow_probability as tfp
import shap
import numpy as np
from sklearn.datasets import make_regression
from sklearn.preprocessing import KBinsDiscretizer
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
N_BINS = 7
INPUT_SIZE = 15<jupyter_output><empty_output><jupyter_text>Make Dataset<jupyter_code># make regression dataset
X, y_c = make_regression(
    n_samples=5000,
    n_features=INPUT_SIZE,
    n_informative=6,
    bias=123.43,
)

# bin the dataset
discretizer = KBinsDiscretizer(
    n_bins=N_BINS,
    encode="onehot-dense",
    strategy="quantile",
)
y = discretizer.fit_transform(y_c.reshape(-1, 1))

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=3)

X_train = X_train.astype("float32")
y_train = y_train.astype("float32")
X_test = X_test.astype("float32")
y_test = y_test.astype("float32")
bin_edges = discretizer.bin_edges_[0].astype("float32")

print(X_train.shape)
print(y_train.shape)
print(X_test.shape)
print(y_test.shape)
print("Bin edges ", bin_edges)<jupyter_output>(4997, 15)
(4997, 7)
(3, 15)
(3, 7)
Bin edges  [-377.4633    -20.754118   43.6177     96.69855   143.89049   199.98409
  268.98087   625.614   ]<jupyter_text>Define a Model<jupyter_code>model = tf.keras.models.Sequential([
  tf.keras.layers.InputLayer(input_shape=(INPUT_SIZE,)),
  tf.keras.layers.Dense(30, activation="relu"),
  tf.keras.layers.Dense(N_BINS, activation="softmax"),
])
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.00001),
    loss=tf.keras.losses.CategoricalCrossentropy(),
)
model.summary()
model.fit(
    x=X_train,
    y=y_train,
    epochs=1,
    batch_size=512,
)<jupyter_output>10/10 [==============================] - 0s 2ms/step - loss: 2.1241<jupyter_text>Computing the Median<jupyter_code>predictions = model(X_test)
quantized_dist = tfp.distributions.MixtureSameFamily(
    mixture_distribution=tfp.distributions.Categorical(probs=predictions[0]),
    components_distribution=tfp.distributions.Uniform(low=bin_edges[:-1], high=bin_edges[1:])
)
samples = quantized_dist.sample(sample_shape=(100000,))
plt.hist(samples, bins=1000)
plt.show()
quantized_dist.cdf(0.0)
#quantized_dist.quantile(0.5)
def sampled_median(probabilities, bin_edges):
  batched_bin_edges = np.broadcast_to(bin_edges, (probabilities.shape[0], N_BINS + 1))
  quantized_dist = tfp.distributions.MixtureSameFamily(
    mixture_distribution=tfp.distributions.Categorical(probs=probabilities),
    components_distribution=tfp.distributions.Uniform(low=batched_bin_edges[:, :-1], high=batched_bin_edges[:, 1:])
  )

  return np.median(quantized_dist.sample(sample_shape=100000).numpy(), axis=0)


sampled_median(model(X_test), bin_edges)
def quantile(probabilities, bin_edges, quantile):
  bin_widths = np.diff(bin_edges)
  bin_densities = probabilities / bin_widths

  bin_quantiles = tf.cumsum(probabilities, axis=1)

  quantile_bin = tf.searchsorted(bin_quantiles, bin_quantiles[:, -1:] * quantile)
  
  lower_edge = tf.gather_nd(bin_edges, quantile_bin)
  higher_edge = tf.gather_nd(bin_edges, quantile_bin + 1)

  bin_upper_quantile = tf.gather_nd(bin_quantiles, quantile_bin, batch_dims=1)
  bin_density = tf.gather_nd(bin_densities, quantile_bin, batch_dims=1)
  bin_width = tf.gather_nd(bin_widths, quantile_bin)

  residual = ((bin_upper_quantile - quantile) /  bin_density)

  return higher_edge - residual

res = quantile(model(X_test), bin_edges, 0.5)
print(res)<jupyter_output>tf.Tensor([102.514275  75.3732   106.69275 ], shape=(3,), dtype=float32)<jupyter_text>Testing Explainers  Add Layer to Compute Median<jupyter_code>model.add(tf.keras.layers.Lambda(lambda probs: quantile(probs, bin_edges, 0.5)))
model.summary()<jupyter_output>Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense (Dense)                (None, 30)                480       
_________________________________________________________________
dense_1 (Dense)              (None, 7)                 217       
_________________________________________________________________
lambda (Lambda)              (None,)                   0         
=================================================================
Total params: 697
Trainable params: 697
Non-trainable params: 0
_________________________________________________________________<jupyter_text>Kernel Explainer<jupyter_code>explainer = shap.KernelExplainer(lambda x: model(x).numpy(), X_train)
shap_values = explainer.shap_values(X_test)
shap.summary_plot(shap_values)
shap.summary_plot(shap_values, plot_type="bar")

shap.initjs()
shap.force_plot(explainer.expected_value, shap_values, X_test)<jupyter_output><empty_output><jupyter_text>Gradient Explainer<jupyter_code>explainer = shap.GradientExplainer(model, X_train)
shap_values = explainer.shap_values(X_test)
shap.summary_plot(shap_values)
shap.summary_plot(shap_values, plot_type="bar")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1387349512707816, 'fraction_numerical': 0.039174469711446586, 'mean_word_length': 5.567126725219573, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '591580', 'license': 'Apache-2.0', 'n_tokens_mistral': 1977, 'n_tokens_neox': 1814, 'n_words': 277, 'path': 'tfp.ipynb', 'repo_name': 'namanwahi/notebooks', 'size': '5233'}