sandbox338 commited on
Commit
faf58d1
·
verified ·
1 Parent(s): bd3e9f0
Files changed (1) hide show
  1. app.py +10 -2
app.py CHANGED
@@ -4,11 +4,19 @@ from transformers import pipeline
4
  # Load the model from Hugging Face Hub
5
  classifier = pipeline("text-classification", model="sandbox338/hatespeech")
6
 
 
 
 
 
 
 
 
7
  # Define the function for the interface
8
  def classify_text(text):
9
  results = classifier(text)
10
- # Return only the label without the score
11
- return results[0]['label']
 
12
 
13
  # Create the interface
14
  interface = gr.Interface(
 
4
  # Load the model from Hugging Face Hub
5
  classifier = pipeline("text-classification", model="sandbox338/hatespeech")
6
 
7
+ # Define the label-to-number mapping
8
+ label_to_number = {
9
+ "Non-hate speech": 0,
10
+ "Political hate speech": 1,
11
+ "Offensive language": 2
12
+ }
13
+
14
  # Define the function for the interface
15
  def classify_text(text):
16
  results = classifier(text)
17
+ # Get the label and map it to the corresponding number
18
+ label = results[0]['label']
19
+ return label_to_number[label]
20
 
21
  # Create the interface
22
  interface = gr.Interface(