[refactor]: replace pyplot usage with matplotlib Figure

app.py

@@ -2,13 +2,13 @@ import random
 from typing import *
 
 import gradio as gr
-import matplotlib.pyplot as plt
 import numpy as np
 import seaborn as sns
 import sentencepiece as sp
 import torch
 
 from huggingface_hub import hf_hub_download
+from matplotlib.figure import Figure
 from torchtext.datasets import Multi30k
 
 from models import Seq2Seq
@@ -31,19 +31,18 @@ normalize = lambda sample: (sample[0].lower().strip(), sample[1].lower().strip()
 test_source, _ = zip(*map(normalize, Multi30k(split="test", language_pair=("de", "en"))))
 
 
-def attention_heatmap(input_tokens: List[str], output_tokens: List[str], weights: np.ndarray) -> plt.Figure:
-    figure = plt.figure(dpi=800, tight_layout=True)
-    axes = sns.heatmap(weights, cmap="gray", cbar=False)
-    axes.set_xticklabels(input_tokens, rotation=90)
-    axes.set_yticklabels(output_tokens, rotation=0)
-    axes.tick_params(axis="both", length=0)
+def attention_heatmap(input_tokens: List[str], output_tokens: List[str], weights: np.ndarray) -> Figure:
+    figure = Figure(dpi=800, tight_layout=True)
+    axes = figure.add_subplot()
+    axes = sns.heatmap(weights, ax=axes, xticklabels=input_tokens, yticklabels=output_tokens, cmap="gray", cbar=False)
+    axes.tick_params(axis="x", rotation=90, length=0)
+    axes.tick_params(axis="y", rotation=0, length=0)
     axes.xaxis.tick_top()
-    plt.close()
     return figure
 
 
 @torch.inference_mode()
-def run(input: str) -> Tuple[str, plt.Figure]:
+def run(input: str) -> Tuple[str, Figure]:
     """Run inference on a single sentence. Returns prediction and attention heatmap."""""
     input = input.lower().strip().rstrip(".") + "."
     input_tensor = torch.tensor(source_spm.encode(input), dtype=torch.int64)