use cache and quantization

app.py

@@ -1,45 +1,31 @@
 from collections import Counter
 
 import graphviz
-from optimum.bettertransformer import BetterTransformer
 import penman
 from penman.models.noop import NoOpModel
-from mbart_amr.constraints.constraints import AMRLogitsProcessor
 from mbart_amr.data.linearization import linearized2penmanstr
-from mbart_amr.data.tokenization import AMRMBartTokenizer
-from transformers import MBartForConditionalGeneration, LogitsProcessorList
+from transformers import LogitsProcessorList
 
 import streamlit as st
 
-if "logits_processor" not in st.session_state:
-    st.session_state["logits_processor"] = None
+from utils import get_resources
 
-if "tokenizer" not in st.session_state:
-    st.session_state["tokenizer"] = None
-
-if "model" not in st.session_state:
-    st.session_state["tokenizer"] = AMRMBartTokenizer.from_pretrained("BramVanroy/mbart-en-to-amr", src_lang="en_XX")
-    st.session_state["model"] = MBartForConditionalGeneration.from_pretrained("BramVanroy/mbart-en-to-amr")
-    st.session_state["model"] = BetterTransformer.transform(st.session_state["model"], keep_original_model=False)
-    st.session_state["model"].resize_token_embeddings(len(st.session_state["tokenizer"]))
-    st.session_state["logits_processor"] = AMRLogitsProcessor(st.session_state["tokenizer"],
-                                                              st.session_state["model"].config.max_length)
+model, tokenizer, logitsprocessor = get_resources()
 
 st.title("📝 Parse text into AMR")
 
 text = st.text_input(label="Text to transform (en)")
 
-if text and "model" in st.session_state:
+if text:
     gen_kwargs = {
-        "max_length": st.session_state["model"].config.max_length,
-        "num_beams": st.session_state["model"].config.num_beams,
-        "logits_processor": LogitsProcessorList([st.session_state["logits_processor"]]) if st.session_state[
-            "logits_processor"] else None
+        "max_length": model.config.max_length,
+        "num_beams": model.config.num_beams,
+        "logits_processor": LogitsProcessorList([logitsprocessor])
     }
 
-    encoded = st.session_state["tokenizer"](text, return_tensors="pt")
-    generated = st.session_state["model"].generate(**encoded, **gen_kwargs)
-    linearized = st.session_state["tokenizer"].decode_and_fix(generated)[0]
+    encoded = tokenizer(text, return_tensors="pt")
+    generated = model.generate(**encoded, **gen_kwargs)
+    linearized = tokenizer.decode_and_fix(generated)[0]
     penman_str = linearized2penmanstr(linearized)
 
     try:

utils.py

@@ -0,0 +1,30 @@
+import streamlit as st
+
+from torch.quantization import quantize_dynamic
+from torch import nn, qint8
+from torch.nn import Parameter
+from transformers import PreTrainedModel, PreTrainedTokenizer
+from optimum.bettertransformer import BetterTransformer
+from mbart_amr.constraints.constraints import AMRLogitsProcessor
+from mbart_amr.data.tokenization import AMRMBartTokenizer
+from transformers import MBartForConditionalGeneration
+
+
+st_hash_funcs = {PreTrainedModel: lambda model: model.name_or_path,
+                 PreTrainedTokenizer: lambda tokenizer: tokenizer.name_or_path,
+                 Parameter: lambda param: param.data}
+
+
+@st.cache(show_spinner=False, hash_funcs=st_hash_funcs, allow_output_mutation=True)
+def get_resources(quantize: bool = True):
+    tokenizer = AMRMBartTokenizer.from_pretrained("BramVanroy/mbart-en-to-amr", src_lang="en_XX")
+    model = MBartForConditionalGeneration.from_pretrained("BramVanroy/mbart-en-to-amr")
+    model = BetterTransformer.transform(model, keep_original_model=False)
+    model.resize_token_embeddings(len(tokenizer))
+
+    if quantize:
+        model = quantize_dynamic(model, {nn.Linear, nn.Dropout, nn.LayerNorm}, dtype=qint8)
+
+    logits_processor = AMRLogitsProcessor(tokenizer, model.config.max_length)
+
+    return model, tokenizer, logits_processor