Cache model and tokenizer and lock dependencies

app.py

@@ -2,10 +2,10 @@ import os
 import sys
 
 import streamlit as st
+import transformers
 from huggingface_hub import snapshot_download
 from transformers import AutoTokenizer
 
-
 LOCAL_PATH = snapshot_download("flax-community/clip-spanish")
 sys.path.append(LOCAL_PATH)
 
@@ -15,16 +15,24 @@ from test_on_image import run_inference
 
 def save_file_to_disk(uplaoded_file):
     temp_file = os.path.join("/tmp", uplaoded_file.name)
-    with open(temp_file,"wb") as f:
+    with open(temp_file, "wb") as f:
         f.write(uploaded_file.getbuffer())
     return temp_file
 
+
+@st.cache(
+    hash_funcs={
+        transformers.models.bert.tokenization_bert_fast.BertTokenizerFast: id,
+        FlaxHybridCLIP: id,
+    }
+)
 def load_tokenizer_and_model():
     # load the saved model
     tokenizer = AutoTokenizer.from_pretrained("dccuchile/bert-base-spanish-wwm-cased")
     model = FlaxHybridCLIP.from_pretrained(LOCAL_PATH)
     return tokenizer, model
 
+
 tokenizer, model = load_tokenizer_and_model()
 
 st.title("Image-Caption Matching")
@@ -36,7 +44,15 @@ if uploaded_file is not None and text_input:
     try:
         local_image_path = save_file_to_disk(uploaded_file)
         score = run_inference(local_image_path, text_input, model, tokenizer).tolist()
-        st.image(uploaded_file, caption=text_input, width=None, use_column_width=None, clamp=False, channels='RGB', output_format='auto')
+        st.image(
+            uploaded_file,
+            caption=text_input,
+            width=None,
+            use_column_width=None,
+            clamp=False,
+            channels="RGB",
+            output_format="auto",
+        )
         st.write(f"## Score: {score:.2f}")
     finally:
         if local_image_path:

requirements.txt

@@ -1,8 +1,8 @@
-flax
-huggingface_hub
-jax
+flax==0.3.4
+huggingface-hub==0.0.12
+jax==0.2.17
 streamlit==0.84.1
-torch
-torchvision
-transformers
-watchdog
+torch==1.9.0
+torchvision==0.10.0
+transformers==4.8.2
+watchdog==2.1.3