Update app.py

app.py

@@ -4,26 +4,24 @@ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from time import time
 import torch
 
-@st.cache_resource(
-    allow_output_mutation=True,
-    hash_funcs={
-        AutoTokenizer: lambda x: None,
-        AutoModelForSeq2SeqLM: lambda x: None,
-    },
-    suppress_st_warning=True
-)
-def load_models(lan):
+
+def load_tok_and_data(lan):
     st_time = time()
     tokenizer = AutoTokenizer.from_pretrained("Babelscape/mrebel-large", src_lang=_Tokens[lan], tgt_lang="tp_XX")
+    dataset = load_dataset('Babelscape/SREDFM', lan, split="validation", streaming=True)
+    dataset = [example for example in dataset.take(1001)]
+    return (tokenizer, dataset)
+    
+@st.cache_resource
+def load_model():
+    st_time = time()
     print("+++++ loading Model", time() - st_time)
     model = AutoModelForSeq2SeqLM.from_pretrained("Babelscape/mrebel-large")
     if torch.cuda.is_available():
         _ = model.to("cuda:0") # comment if no GPU available
     _ = model.eval()
     print("+++++ loaded model", time() - st_time)
-    dataset = load_dataset('Babelscape/SREDFM', lan, split="validation", streaming=True)
-    dataset = [example for example in dataset.take(1001)]
-    return (tokenizer, model, dataset)
+    return model
 
 def extract_triplets_typed(text):
     triplets = []
@@ -63,13 +61,15 @@ def extract_triplets_typed(text):
 
 st.markdown("""This is a demo for the Findings of EMNLP 2021 paper [REBEL: Relation Extraction By End-to-end Language generation](https://github.com/Babelscape/rebel/blob/main/docs/EMNLP_2021_REBEL__Camera_Ready_.pdf). The pre-trained model is able to extract triplets for up to 200 relation types from Wikidata or be used in downstream Relation Extraction task by fine-tuning. Find the model card [here](https://huggingface.co/Babelscape/rebel-large). Read more about it in the [paper](https://aclanthology.org/2021.findings-emnlp.204) and in the original [repository](https://github.com/Babelscape/rebel).""")
 
+model = load_model()
+
 lan = st.selectbox(
     'Select a Language',
     ('ar', 'ca', 'de', 'el', 'en', 'es', 'fr', 'hi', 'it', 'ja', 'ko', 'nl', 'pl', 'pt', 'ru', 'sv', 'vi', 'zh'), index=1)
 
 _Tokens = {'en': 'en_XX', 'de': 'de_DE', 'ca': 'ca_XX', 'ar': 'ar_AR', 'el': 'el_EL', 'it': 'it_IT', 'ja': 'ja_XX', 'ko': 'ko_KR', 'hi': 'hi_IN', 'pt': 'pt_XX', 'ru': 'ru_RU', 'pl': 'pl_PL', 'zh': 'zh_CN', 'fr': 'fr_XX', 'vi': 'vi_VN', 'sv':'sv_SE'}
 
-tokenizer, model, dataset = load_models(lan)
+tokenizer, dataset = load_tok_and_data(lan)
 
 agree = st.checkbox('Free input', False)
 if agree: