add gpu support

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed
 from transformers import pipeline
+import torch
 import json
 
 @st.cache(allow_output_mutation=True)
@@ -20,10 +21,8 @@ def load_examples():
 
 st.set_page_config(page_icon=':parrot:', layout="wide")
 
-default_code = '''\
-def print_hello_world():\
-'''
-
+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)
@@ -36,7 +35,7 @@ gen_kwargs = {}
 st.title("CodeParrot 🦜")
 st.markdown('##')
 
-pipe = pipeline('text-generation', model=model, tokenizer=tokenizer)
+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"]
@@ -48,7 +47,7 @@ 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..."):