Update app.py

app.py

@@ -15,17 +15,21 @@ st.write("This tool lets you extract relation triples concerning interactions be
 st.write("It is the result of an end of studies project within ESI school and dedicated to biomedical researchers looking to extract precise information about the subject without digging into long publications.")
 
 
-@st.cache
-def load_models():
-  tokenizer = AutoTokenizer.from_pretrained("dmis-lab/biobert-large-cased-v1.1", truncation = True, padding=True, model_max_length=512,)
+@st.cache(allow_output_mutation = True)
+def load_tokenizer():
+  return AutoTokenizer.from_pretrained("dmis-lab/biobert-large-cased-v1.1", truncation = True, padding=True, model_max_length=512,)
+  
+tokenizer = load_tokenizer()      
+
+@st.cache(allow_output_mutation = True)
+def load_modelNER(tokenizer):
   model_checkpoint = BertForTokenClassification.from_pretrained("dexay/Ner2HgF", )
-   
+  return pipeline("token-classification", tokenizer = tokenizer,model=model_checkpoint,  )
+@st.cache(allow_output_mutation = True)
+def load_modelRE(tokenizer):
   model_re = AutoModelForSequenceClassification.from_pretrained("dexay/reDs3others", )
+  return pipeline("text-classification", tokenizer = tokenizer,model=model_re, )
   
-  return tokenizer , model_checkpoint , model_re
-
-
-
 form = st.form(key='my-form')
 x = form.text_area('Enter text', height=250)
 submit = form.form_submit_button('Submit')
@@ -38,14 +42,10 @@ if submit and len(x) != 0:
   #model.to("cpu")
   st.text("Execution is in progress ...")
   
-  tmr = load_models()
-  
-  tokenizer = tmr[0]
-  model_checkpoint = tmr[1]
-  model_re = tmr[2]
+
   
   
-  token_classifier = pipeline("token-classification", tokenizer = tokenizer,model=model_checkpoint,  )
+  #token_classifier = pipeline("token-classification", tokenizer = tokenizer,model=model_checkpoint,  )
   
   
   
@@ -74,6 +74,7 @@ if submit and len(x) != 0:
   
   #tokenized_dat = tokenize_function(ddata) 
   
+  token_classifier = load_modelNER(tokenizer)
   az = token_classifier(ddata)
   
   
@@ -189,13 +190,12 @@ if submit and len(x) != 0:
   
   # Relation extraction part
   
-  token_classifier = pipeline("text-classification", tokenizer = tokenizer,model=model_re, 
-  )
+  #token_classifier = pipeline("text-classification", tokenizer = tokenizer,model=model_re, )
   
   rrdata = lstSentEnc
   
   
-  
+  token_classifier = load_modelRE(tokenizer)
   outre = token_classifier(rrdata)