jayebaku commited on
Commit
01130ab
·
verified ·
1 Parent(s): 6f7ff99

Create classifier.py

Browse files
Files changed (1) hide show
  1. classifier.py +17 -0
classifier.py ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline
2
+
3
+ def classify(tweet, event_model, hftoken):
4
+
5
+ # event type prediction
6
+ event_predictor = pipeline(task="text-classification", model=event_model,
7
+ batch_size=512, token=hftoken)
8
+ tokenizer_kwargs = {'padding': True, 'truncation': True, 'max_length': 512}
9
+
10
+ results = {"text": None, "event": None, "score": None}
11
+ prediction = event_predictor(tweet, **tokenizer_kwargs)[0]
12
+
13
+ results["text"] = tweet
14
+ results["event"] = prediction["label"]
15
+ results["score"] = prediction["score"]
16
+
17
+ return results