Update app.py

app.py

@@ -4,7 +4,6 @@ import pandas as pd
 import os
 import torch
 import torch.nn as nn
-from transformers import ElectraModel, AutoConfig, GPT2LMHeadModel
 from transformers.activations import get_activation
 from transformers import AutoTokenizer, AutoModelWithLMHead
 
@@ -13,10 +12,14 @@ st.title('GPT2 ')
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-
-
-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")
+    model.to(device)
+    return model, tokenizer
+    
+model, tokenizer = get_model()
 
 g = """informal english: garage band has made people who know nothing about music good at creating music.
 Translated into the Style of Abraham Lincoln: garage band ( offers the uninitiated in music the ability to produce professional-quality compositions / catapults those for whom music is an uncharted art the ability the realize masterpieces / stimulates music novice's competency to yield sublime arrangements / begets individuals of rudimentary musical talent the proficiency to fashion elaborate suites ).