Add code to select the number of GPU layers

app.py

@@ -1,6 +1,6 @@
+#!/usr/bin/env python3
 import gradio as gr
 import logging
-import re
 from llama_cpp import Llama
 from huggingface_hub import hf_hub_download
 import sys
@@ -46,6 +46,7 @@ def generate_text(prompt, example):
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--model", help="Path to the model file")
+parser.add_argument("--gpu", help="How many GPU layers to use", default=0, type=int)
 args = parser.parse_args()
 
 if args.model:
@@ -53,7 +54,8 @@ if args.model:
 else:
     model_path = download_model()
 
-llm = Llama(model_path=model_path)
+# TODO n_gpu_layers doesn't seem to be working
+llm = Llama(model_path=model_path, n_gpu_layers=args.gpu)
 
 # Create a Gradio interface
 interface = gr.Interface(