leeex00 commited on
Commit
fdd57e1
·
verified ·
1 Parent(s): cc85e9e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -0
app.py CHANGED
@@ -1,4 +1,5 @@
1
  import gradio as gr
 
2
  from transformers import AutoTokenizer, AutoModelForCausalLM
3
  import torch
4
 
@@ -12,9 +13,11 @@ model = AutoModelForCausalLM.from_pretrained(
12
  trust_remote_code=True
13
  )
14
 
 
15
  def generate(prompt):
16
  inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
17
  outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.7)
18
  return tokenizer.decode(outputs[0], skip_special_tokens=True)
19
 
20
  gr.Interface(fn=generate, inputs="text", outputs="text", title="WizardCoder 13B").launch()
 
 
1
  import gradio as gr
2
+ import spaces
3
  from transformers import AutoTokenizer, AutoModelForCausalLM
4
  import torch
5
 
 
13
  trust_remote_code=True
14
  )
15
 
16
+ @spaces.GPU
17
  def generate(prompt):
18
  inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
19
  outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.7)
20
  return tokenizer.decode(outputs[0], skip_special_tokens=True)
21
 
22
  gr.Interface(fn=generate, inputs="text", outputs="text", title="WizardCoder 13B").launch()
23
+