jayebaku commited on
Commit
94a7f16
·
verified ·
1 Parent(s): b223991

Update classifier.py

Browse files
Files changed (1) hide show
  1. classifier.py +2 -2
classifier.py CHANGED
@@ -1,12 +1,12 @@
1
  import spaces
2
  from transformers import pipeline
3
 
4
- @spaces.GPU(duration=60)
5
  def classify(tweet, event_model, hftoken, threshold):
6
 
7
  # event type prediction
8
  event_predictor = pipeline(task="text-classification", model=event_model,
9
- batch_size=512, token=hftoken)
10
  tokenizer_kwargs = {'padding': True, 'truncation': True, 'max_length': 512}
11
 
12
  results = {"text": None, "event": None, "score": None}
 
1
  import spaces
2
  from transformers import pipeline
3
 
4
+ #@spaces.GPU(duration=60)
5
  def classify(tweet, event_model, hftoken, threshold):
6
 
7
  # event type prediction
8
  event_predictor = pipeline(task="text-classification", model=event_model,
9
+ batch_size=512, token=hftoken, device_map="auto")
10
  tokenizer_kwargs = {'padding': True, 'truncation': True, 'max_length': 512}
11
 
12
  results = {"text": None, "event": None, "score": None}