philipp-zettl commited on
Commit
754acee
·
verified ·
1 Parent(s): d8ded2a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -1,16 +1,18 @@
1
  import gradio as gr
2
  from transformers import AutoTokenizer
3
  from model import MultiHeadClassification
 
4
 
5
 
6
  model = MultiHeadClassification.from_pretrained(
7
  'philipp-zettl/multi-head-sequence-classification-model',
8
  {"GGU": 3, "sentiment": 3}
9
- )
10
  tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-m3')
11
 
 
12
  def generate(prompt):
13
- inputs = tokenizer([prompt])
14
  return model(**inputs)
15
 
16
 
 
1
  import gradio as gr
2
  from transformers import AutoTokenizer
3
  from model import MultiHeadClassification
4
+ import spaces
5
 
6
 
7
  model = MultiHeadClassification.from_pretrained(
8
  'philipp-zettl/multi-head-sequence-classification-model',
9
  {"GGU": 3, "sentiment": 3}
10
+ ).to('cuda')
11
  tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-m3')
12
 
13
+ @spaces.GPU()
14
  def generate(prompt):
15
+ inputs = tokenizer([prompt]).to('cuda')
16
  return model(**inputs)
17
 
18