Command line argument for context size

app.py

@@ -47,6 +47,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)
+parser.add_argument("--ctx", help="How many context tokens to use", default=512, type=int)
 args = parser.parse_args()
 
 if args.model:
@@ -55,7 +56,7 @@ else:
     model_path = download_model()
 
 # TODO n_gpu_layers doesn't seem to be working
-llm = Llama(model_path=model_path, n_gpu_layers=args.gpu)
+llm = Llama(model_path=model_path, n_gpu_layers=args.gpu, n_ctx=args.ctx)
 
 # Create a Gradio interface
 interface = gr.Interface(