Update mT5Model.py

mT5Model.py

@@ -1,10 +1,12 @@
 from torch.nn.functional import softmax
 from transformers import MT5ForConditionalGeneration, MT5Tokenizer
+import streamlit as st
 
 def process_nli(premise: str, hypothesis: str):
     """ process to required xnli format with task prefix """
     return "".join(['xnli: premise: ', premise, ' hypothesis: ', hypothesis])
 
+@st.cache(allow_output_mutation=True)
 def setModel(model_name):
     tokenizer = MT5Tokenizer.from_pretrained(model_name)
     model = MT5ForConditionalGeneration.from_pretrained(model_name)