Mismatch num_samples for Torch and Tensorflow implementation IG about path_explain HOT 7 CLOSED

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from path_explain import PathExplainerTorch, text_plot

model = AutoModelForCausalLM.from_pretrained("distilgpt2")
tokenizer = AutoTokenizer.from_pretrained("distilgpt2")

def embedding_model(batch_ids):
    batch_embedding = model.transformer.wte.weight[[batch_ids]]
    return batch_embedding

def prediction_model(batch_embedding):
    logits = model(inputs_embeds=batch_embedding).logits[:, -1]
    probs = torch.log_softmax(logits, dim=-1)
    
    return probs

sen = "It is raining cats and"

batch_ids = tokenizer([sen])['input_ids']
batch_embedding = embedding_model(batch_ids)
baseline_embedding = torch.zeros_like(batch_embedding)
output_idx = torch.tensor([tokenizer.convert_tokens_to_ids(f"Ġdogs")])

explainer = PathExplainerTorch(prediction_model)

attributions = explainer.attributions(
    batch_embedding,
    baseline_embedding,
    output_indices=output_idx,
    num_samples=9,
    use_expectation=False,
)

sum_attr = attributions.squeeze().detach().numpy().sum(-1)  # Sum over embedding dim

text_plot(
    sen.split(),
    sum_attr,
)

import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf

from scipy.special import log_softmax
from transformers import TFAutoModelForCausalLM, AutoTokenizer

from path_explain import EmbeddingExplainerTF, text_plot

tf_model = TFAutoModelForCausalLM.from_pretrained("distilgpt2")
tf_tokenizer = AutoTokenizer.from_pretrained("distilgpt2")

def embedding_model(batch_ids):
    batch_embedding = tf_model.transformer.wte(batch_ids)
    return batch_embedding

def prediction_model(batch_embedding):
    if isinstance(batch_embedding, np.ndarray):
        batch_embedding = tf.convert_to_tensor(batch_embedding)
    logits = tf_model(None, inputs_embeds=batch_embedding).logits[:, -1]
    probs = tf.nn.log_softmax(logits, axis=-1)
    
    return probs

sen = "It is raining cats and"

batch_ids = tf_tokenizer([sen])['input_ids']
batch_embedding = embedding_model(batch_ids)

baseline_embedding = np.zeros_like(batch_embedding)

output_idx = tf_tokenizer.convert_tokens_to_ids("Ġdogs")

explainer = EmbeddingExplainerTF(prediction_model)
attributions = explainer.attributions(
    batch_embedding,
    baseline_embedding,
    output_indices=output_idx,
    num_samples=10,
    batch_size=1,
    use_expectation=False,
)

text_plot(
    sen.split(),
    attributions[0],
)