filipealmeida
commited on
Commit
•
18d0ebf
1
Parent(s):
2c40787
Add code to select the number of GPU layers
Browse files
app.py
CHANGED
@@ -1,6 +1,6 @@
|
|
|
|
1 |
import gradio as gr
|
2 |
import logging
|
3 |
-
import re
|
4 |
from llama_cpp import Llama
|
5 |
from huggingface_hub import hf_hub_download
|
6 |
import sys
|
@@ -46,6 +46,7 @@ def generate_text(prompt, example):
|
|
46 |
|
47 |
parser = argparse.ArgumentParser()
|
48 |
parser.add_argument("--model", help="Path to the model file")
|
|
|
49 |
args = parser.parse_args()
|
50 |
|
51 |
if args.model:
|
@@ -53,7 +54,8 @@ if args.model:
|
|
53 |
else:
|
54 |
model_path = download_model()
|
55 |
|
56 |
-
|
|
|
57 |
|
58 |
# Create a Gradio interface
|
59 |
interface = gr.Interface(
|
|
|
1 |
+
#!/usr/bin/env python3
|
2 |
import gradio as gr
|
3 |
import logging
|
|
|
4 |
from llama_cpp import Llama
|
5 |
from huggingface_hub import hf_hub_download
|
6 |
import sys
|
|
|
46 |
|
47 |
parser = argparse.ArgumentParser()
|
48 |
parser.add_argument("--model", help="Path to the model file")
|
49 |
+
parser.add_argument("--gpu", help="How many GPU layers to use", default=0, type=int)
|
50 |
args = parser.parse_args()
|
51 |
|
52 |
if args.model:
|
|
|
54 |
else:
|
55 |
model_path = download_model()
|
56 |
|
57 |
+
# TODO n_gpu_layers doesn't seem to be working
|
58 |
+
llm = Llama(model_path=model_path, n_gpu_layers=args.gpu)
|
59 |
|
60 |
# Create a Gradio interface
|
61 |
interface = gr.Interface(
|