Upload app.py
Browse files
app.py
CHANGED
@@ -14,13 +14,14 @@ tokenizer = AutoTokenizer.from_pretrained(model_path)
|
|
14 |
|
15 |
model = AutoModelForCausalLM.from_pretrained( model_path)
|
16 |
|
|
|
17 |
@spaces.GPU
|
18 |
def generate_text(prompt, temperature=0.9, max_length=1200):
|
19 |
# Encode the inputs
|
20 |
inputs = tokenizer.encode(prompt, return_tensors="pt")
|
21 |
attention_mask = torch.ones(inputs.shape, dtype=torch.long)
|
22 |
-
inputs = inputs.to(
|
23 |
-
attention_mask = attention_mask.to(
|
24 |
outputs = model.generate(
|
25 |
inputs,
|
26 |
attention_mask=attention_mask,
|
@@ -35,7 +36,7 @@ def generate_text(prompt, temperature=0.9, max_length=1200):
|
|
35 |
def gradio_app():
|
36 |
with gr.Blocks() as demo:
|
37 |
gr.Markdown(title)
|
38 |
-
prompt = gr.Code(label="Enter your code prompt", value="def
|
39 |
with gr.Row():
|
40 |
temperature = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="Temperature")
|
41 |
max_length = gr.Slider(minimum=100, maximum=1024, step=10, value=100, label="Generate Length")
|
|
|
14 |
|
15 |
model = AutoModelForCausalLM.from_pretrained( model_path)
|
16 |
|
17 |
+
device = 'cpu' # 'cuda'
|
18 |
@spaces.GPU
|
19 |
def generate_text(prompt, temperature=0.9, max_length=1200):
|
20 |
# Encode the inputs
|
21 |
inputs = tokenizer.encode(prompt, return_tensors="pt")
|
22 |
attention_mask = torch.ones(inputs.shape, dtype=torch.long)
|
23 |
+
inputs = inputs.to(device)
|
24 |
+
attention_mask = attention_mask.to(device)
|
25 |
outputs = model.generate(
|
26 |
inputs,
|
27 |
attention_mask=attention_mask,
|
|
|
36 |
def gradio_app():
|
37 |
with gr.Blocks() as demo:
|
38 |
gr.Markdown(title)
|
39 |
+
prompt = gr.Code(label="Enter your code prompt", value="def prime_number(n_max):")
|
40 |
with gr.Row():
|
41 |
temperature = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="Temperature")
|
42 |
max_length = gr.Slider(minimum=100, maximum=1024, step=10, value=100, label="Generate Length")
|