mrfakename commited on
Commit
72844d2
·
1 Parent(s): ab49618

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -0
app.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import torch
2
  from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
3
  import gradio as gr
@@ -18,6 +19,8 @@ model = AutoModelForCausalLM.from_pretrained(
18
  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
19
  trust_remote_code=True,
20
  ).to(device)
 
 
21
  def generate_text(text, temperature, maxLen):
22
  inputs = tokenizer([text], return_tensors="pt").to(device)
23
  streamer = TextIteratorStreamer(tokenizer)
 
1
+ import spaces
2
  import torch
3
  from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
4
  import gradio as gr
 
19
  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
20
  trust_remote_code=True,
21
  ).to(device)
22
+
23
+ @spaces.GPU
24
  def generate_text(text, temperature, maxLen):
25
  inputs = tokenizer([text], return_tensors="pt").to(device)
26
  streamer = TextIteratorStreamer(tokenizer)