Brian-A commited on
Commit
e2e3ac7
Β·
verified Β·
1 Parent(s): e8ede39
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import gradio as gr
2
  import torch
3
  from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
@@ -15,6 +16,7 @@ class StopOnTokens(StoppingCriteria):
15
  return True
16
  return False
17
 
 
18
  def predict(message, history):
19
  history_transformer_format = list(zip(history[:-1], history[1:])) + [[message, ""]]
20
  stop = StopOnTokens()
 
1
+ import spaces
2
  import gradio as gr
3
  import torch
4
  from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
 
16
  return True
17
  return False
18
 
19
+ @spaces.GPU
20
  def predict(message, history):
21
  history_transformer_format = list(zip(history[:-1], history[1:])) + [[message, ""]]
22
  stop = StopOnTokens()