Update app.py

app.py

@@ -1,10 +1,7 @@
-import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed, pipeline
 
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
-
-title = "Santacoder 🎅 bash/shell 🐚 Completion"
+title = "SantaCoder 🎅 bash/shell 🐚 Completion"
 description = "This is a subspace to make code generation with [SantaCoder fine-tuned on The Stack bash/shell](https://huggingface.co/mrm8488/santacoder-finetuned-the-stack-bash-4)"
 EXAMPLE_0 = "#!/bin/bash\n# This script removes files larger than 2MB in the current folder\nfind ."
 EXAMPLE_1 = "#!/bin/bash\n\n# This script send an email\nto=”[email protected]”\nsubject=”Greeting”\nmsg=”Welcome to our site”\n"
@@ -18,7 +15,7 @@ model = AutoModelForCausalLM.from_pretrained("mrm8488/santacoder-finetuned-the-s
 
 def code_generation(gen_prompt, max_tokens, temperature=0.6, seed=42):
     set_seed(seed)
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
+    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)
     generated_text = pipe(gen_prompt, do_sample=True, top_p=0.95, temperature=temperature, max_new_tokens=max_tokens)[0]['generated_text']
     return generated_text