modify loading logic to match highlighting

app.py

@@ -5,10 +5,14 @@ import torch
 import json
 
 
-#@st.cache(allow_output_mutation=True)
-def load_pipeline(model_ckpt, device):
-    pipe = pipeline('text-generation', model=model_ckpt, device=device)
-    return pipe
+@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)
+    return model
 
 @st.cache()
 def load_examples():
@@ -18,10 +22,11 @@ def load_examples():
 
 st.set_page_config(page_icon=':parrot:', layout="wide")
 
-device = 0 if torch.cuda.is_available() else -1
-device_name = "GPU" if torch.cuda.is_available() else "CPU"
 model_ckpt = "lvwerra/codeparrot"
-pipe = load_pipeline(model_ckpt, device)
+tokenizer = load_tokenizer(model_ckpt)
+model = load_model(model_ckpt)
+pipe = pipeline(model=model, tokenizer=tokenizer)
+
 examples = load_examples()
 example_names = [example["name"] for example in examples]
 name2id = dict([(name, i) for i, name in enumerate(example_names)])
@@ -42,7 +47,6 @@ if gen_kwargs["do_sample"]:
     gen_kwargs["temperature"] = st.sidebar.slider("Temperature", value = 0.2, min_value = 0.0, max_value=2.0, step=0.05)
     gen_kwargs["top_k"] = st.sidebar.slider("Top-k", min_value = 0, max_value=100, value = 0)
     gen_kwargs["top_p"] = st.sidebar.slider("Top-p", min_value = 0.0, max_value=1.0, step = 0.01, value = 0.95)
-st.sidebar.markdown(f"Device: _{device_name}_")
 gen_prompt = st.text_area("Generate code with prompt:", value=example_text, height=220,).strip()
 if st.button("Generate code!"):
     with st.spinner("Generating code..."):