Upload app.py

app.py

@@ -14,13 +14,14 @@ tokenizer = AutoTokenizer.from_pretrained(model_path)
 
 model = AutoModelForCausalLM.from_pretrained( model_path)
 
+device = 'cpu' # 'cuda'
 @spaces.GPU
 def generate_text(prompt, temperature=0.9, max_length=1200):
     # Encode the inputs
     inputs = tokenizer.encode(prompt, return_tensors="pt")
     attention_mask = torch.ones(inputs.shape, dtype=torch.long)
-    inputs = inputs.to("cuda")
-    attention_mask = attention_mask.to("cuda")    
+    inputs = inputs.to(device)
+    attention_mask = attention_mask.to(device)    
     outputs = model.generate(
         inputs, 
         attention_mask=attention_mask,
@@ -35,7 +36,7 @@ def generate_text(prompt, temperature=0.9, max_length=1200):
 def gradio_app():
     with gr.Blocks() as demo:
         gr.Markdown(title)
-        prompt = gr.Code(label="Enter your code prompt", value="def print_hello_world():")
+        prompt = gr.Code(label="Enter your code prompt", value="def prime_number(n_max):")
         with gr.Row():
             temperature = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="Temperature")
             max_length = gr.Slider(minimum=100, maximum=1024, step=10, value=100, label="Generate Length")