Remove offensive speech model

app.py

@@ -6,31 +6,23 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 hate_model = AutoModelForSequenceClassification.from_pretrained("KoalaAI/HateSpeechDetector")
 hate_tokenizer = AutoTokenizer.from_pretrained("KoalaAI/HateSpeechDetector")
 
-offensive_model = AutoModelForSequenceClassification.from_pretrained("KoalaAI/OffensiveSpeechDetector")
-offensive_tokenizer = AutoTokenizer.from_pretrained("KoalaAI/OffensiveSpeechDetector")
-
 # Define a function that takes an input text and returns the scores from the models
 def get_scores(text):
   # Tokenize and encode the input text
   hate_input = hate_tokenizer(text, return_tensors="pt")
-  offensive_input = offensive_tokenizer(text, return_tensors="pt")
 
   # Get the logits from the models
   hate_logits = hate_model(**hate_input).logits
-  offensive_logits = offensive_model(**offensive_input).logits
 
   # Apply softmax to get probabilities
   hate_probs = hate_logits.softmax(dim=1)
-  offensive_probs = offensive_logits.softmax(dim=1)
 
   # Get the labels from the models
   hate_labels = hate_model.config.id2label
-  offensive_labels = offensive_model.config.id2label
 
   # Format the output as a dictionary of scores
   output = {}
   output["Hate speech"] = {hate_labels[i]: round(p.item(), 4) for i, p in enumerate(hate_probs[0])}
-  output["Offensive speech"] = {offensive_labels[i]: round(p.item(), 4) for i, p in enumerate(offensive_probs[0])}
 
   return output