Store images instead of showing in a server

I am running the code found on this site in my server and I would like to store images instead of showing them since I have connected remotely with an ssh connection to my server via an SSH connection.

The code is for instance this one:

skip_tokens = [1]  # skip the special token for the start of the text <s>
inp = TextTokenInput(
  eval_prompt, 
  tokenizer,
  skip_tokens=skip_tokens,
)

target = "playing guitar, hiking, and spending time with his family."
attr_res = llm_attr.attribute(inp, target=target, skip_tokens=skip_tokens)
attr_res.plot_token_attr(show=True)

How to store the files locally instead of showing them?

I can't test it but ...

I checked source code and it uses matplotlib for this.

If you remove show=True then it shouldn't show it but it should only get fig, ax.

I think you could use matplotlib.pyplot.savefig(filename) to save it in file.

import matplotlib.pyplot as plt

# ... code  ...

attr_res.plot_token_attr()  # without `show=True
plt.savefig("output.png")
#plt.show()  # eventually show it after saving

Probably you can also use fig for this

fig, ax = attr_res.plot_token_attr()  # without `show=True
fig.savefig("output.png")