Update app.py

app.py

@@ -4,9 +4,8 @@ import pandas as pd
 import os
 import torch
 import torch.nn as nn
-from transformers import ElectraModel, AutoConfig, GPT2LMHeadModel
+from transformers import AutoTokenizer, AutoModelWithLMHead
 from transformers.activations import get_activation
-from transformers import AutoTokenizer
 
 
 st.title('Informal to Formal')
@@ -35,10 +34,13 @@ Translated into the Style of Abraham Lincoln: corn fields ( permeate illinois /
 
 informal english:''')
 
-
-from transformers import AutoTokenizer, AutoModelWithLMHead
-tokenizer = AutoTokenizer.from_pretrained("gpt2")
-model = AutoModelWithLMHead.from_pretrained("BigSalmon/MrLincoln12")
+@st.cache(allow_output_mutation=True)
+def get_model():
+    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    model = AutoModelWithLMHead.from_pretrained("BigSalmon/MrLincoln12")
+    return model, tokenizer
+    
+model, tokenizer = get_model()
 
 with st.form(key='my_form'):
     prompt = st.text_area(label='Enter sentence')