Add cache

generator.py

@@ -26,6 +26,7 @@ import numpy as np
 from transformers import pipeline
 from transformers.pipelines import AggregationStrategy
 import pickle
+import streampit as st
 
 # sq_tokenizer = att.from_pretrained("mrm8488/t5-base-finetuned-question-generation-ap")
 # sq_model = alwm.from_pretrained("mrm8488/t5-base-finetuned-question-generation-ap")
@@ -35,6 +36,7 @@ import pickle
 hfmodel = alwm.from_pretrained("ThomasSimonini/t5-end2end-question-generation")
 
 hftokenizer = T5TokenizerFast.from_pretrained("t5-base")
+@st.cache(suppress_st_warning=True)
 def run_model(input_string, **generator_args):
   generator_args = {
   "max_length": 256,
@@ -61,6 +63,7 @@ tokenizer = att.from_pretrained("ahotrod/albert_xxlargev1_squad2_512")
 model = amqa.from_pretrained("ahotrod/albert_xxlargev1_squad2_512")
 # al_model = pickle.load(open('models/al_model.sav', 'rb'))
 # al_tokenizer = pickle.load(open('models/al_tokenizer.sav', 'rb'))
+@st.cache
 def QA(question, context):
   # model_name="deepset/electra-base-squad2"
   # nlp = pipeline("question-answering",model=al_model,tokenizer=al_tokenizer)