Allow clearing of cache

app.py

@@ -1,10 +1,9 @@
 from io import BytesIO
 import streamlit as st
 import pandas as pd
-import json
 import os
 import numpy as np
-from streamlit.elements import markdown
+from streamlit import caching
 from PIL import Image
 from model.flax_clip_vision_marian.modeling_clip_vision_marian import (
     FlaxCLIPVisionMarianMT,
@@ -31,7 +30,7 @@ tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-es")
 
 @st.cache(persist=True)
 def generate_sequence(pixel_values, num_beams, temperature, top_p):
-    output_ids = model.generate(input_ids=pixel_values, max_length=64, num_beams=num_beams, temperature=temperature, top_p = top_p)
+    output_ids = state.model.generate(input_ids=pixel_values, max_length=64, num_beams=num_beams, temperature=temperature, top_p = top_p)
     print(output_ids)
     output_sequence = tokenizer.batch_decode(output_ids[0], skip_special_tokens=True, max_length=64)
     return output_sequence
@@ -60,7 +59,8 @@ st.sidebar.title("Generation Parameters")
 num_beams = st.sidebar.number_input("Number of Beams", min_value=2, max_value=10, value=4, step=1, help="Number of beams to be used in beam search.")
 temperature = st.sidebar.select_slider("Temperature", options = list(np.arange(0.0,1.1, step=0.1)), value=1.0, help ="The value used to module the next token probabilities.", format_func=lambda x: f"{x:.2f}")
 top_p = st.sidebar.select_slider("Top-P", options = list(np.arange(0.0,1.1, step=0.1)),value=1.0, help="Nucleus Sampling : If set to float < 1, only the most probable tokens with probabilities that add up to :obj:`top_p` or higher are kept for generation.", format_func=lambda x: f"{x:.2f}")
-
+if st.sidebar.button("Clear All Cache"):
+    caching.clear_cache()
 
 image_col, intro_col = st.beta_columns([3, 8])
 image_col.image("./misc/sic-logo.png", use_column_width="always")
@@ -84,6 +84,10 @@ with st.beta_expander("Article"):
     st.write(read_markdown("acknowledgements.md"))
 
 
+if state.model is None:
+    with st.spinner("Loading model..."):
+        state.model = load_model(checkpoints[0])
+
 first_index = 20
 # Init Session State
 if state.image_file is None:
@@ -124,8 +128,7 @@ new_col2.markdown(
     f"""**English Translation**: {translate(state.caption, 'en')}"""
 )
 
-with st.spinner("Loading model..."):
-    model = load_model(checkpoints[0])
+
 sequence = ['']
 if new_col2.button("Generate Caption", help="Generate a caption in the specified language."):
     with st.spinner("Generating Sequence..."):