refactor pipeline loading

app.py

@@ -4,13 +4,10 @@ from transformers import pipeline
 import torch
 import json
 
-@st.cache(allow_output_mutation=True)
-def load_tokenizer(model_ckpt):
-    return AutoTokenizer.from_pretrained(model_ckpt)
 
 @st.cache(allow_output_mutation=True)
-def load_model(model_ckpt):
-    model = AutoModelForCausalLM.from_pretrained(model_ckpt)
+def load_pipeline(model_ckpt, device):
+    pipe = pipeline('text-generation', model=model_ckpt, device=device)
     return model
 
 @st.cache()
@@ -24,8 +21,7 @@ st.set_page_config(page_icon=':parrot:', layout="wide")
 device = 1 if torch.cuda.is_available() else 0
 device_name = "GPU" if torch.cuda.is_available() else "CPU"
 model_ckpt = "lvwerra/codeparrot"
-tokenizer = load_tokenizer(model_ckpt)
-model = load_model(model_ckpt)
+pipe = load_pipe(model_ckpt, device)
 examples = load_examples()
 example_names = [example["name"] for example in examples]
 name2id = dict([(name, i) for i, name in enumerate(example_names)])
@@ -35,7 +31,6 @@ gen_kwargs = {}
 st.title("CodeParrot 🦜")
 st.markdown('##')
 
-pipe = pipeline('text-generation', model=model, tokenizer=tokenizer, device=device)
 st.sidebar.header("Examples:")
 selected_example = st.sidebar.selectbox("Select one of the following examples:", example_names)
 example_text = examples[name2id[selected_example]]["value"]