Xylor commited on
Commit
0c97a7e
·
verified ·
1 Parent(s): b87d52c

Use BERT Zero-Shot Classifier

Browse files
Files changed (1) hide show
  1. app.py +27 -1
app.py CHANGED
@@ -1,10 +1,36 @@
1
  import gradio
2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3
  def my_inference_function(name):
4
  return "Hello " + name + "!"
5
 
6
  gradio_interface = gradio.Interface(
7
- fn = my_inference_function,
 
8
  inputs = "text",
9
  outputs = "text"
10
  )
 
1
  import gradio
2
 
3
+ from transformers import pipeline
4
+ classifier = pipeline("zero-shot-classification",
5
+ model="facebook/bart-large-mnli")
6
+
7
+ # sequence_to_classify = "one day I will see the world"
8
+ # candidate_labels = ['travel', 'cooking', 'dancing']
9
+ # CATEGORIES = ['doc_type.jur', 'doc_type.Spec', 'doc_type.ZDF', 'doc_type.Publ',
10
+ # 'doc_type.Scheme', 'content_type.Alt', 'content_type.Krypto',
11
+ # 'content_type.Karte', 'content_type.Banking', 'content_type.Reg',
12
+ # 'content_type.Konto']
13
+ categories = [
14
+ "Legal", "Specification", "Facts and Figures",
15
+ "Publication", "Payment Scheme",
16
+ "Alternative Payment Systems", "Crypto Payments",
17
+ "Card Payments", "Banking", "Regulations", "Account Payments"
18
+ ]
19
+
20
+ def clf_text(txt: str):
21
+ return classifier(txt, categories)
22
+ # classifier(sequence_to_classify, candidate_labels)
23
+ #{'labels': ['travel', 'dancing', 'cooking'],
24
+ # 'scores': [0.9938651323318481, 0.0032737774308770895, 0.002861034357920289],
25
+ # 'sequence': 'one day I will see the world'}
26
+
27
+
28
  def my_inference_function(name):
29
  return "Hello " + name + "!"
30
 
31
  gradio_interface = gradio.Interface(
32
+ # fn = my_inference_function,
33
+ fn = clf_text,
34
  inputs = "text",
35
  outputs = "text"
36
  )