App 2.8

app.py

@@ -31,7 +31,8 @@ def get_highest_toxicity_class(prediction):
     return model.config.id2label[max_index], prediction[max_index]
 
 input = tokenizer(text, return_tensors="tf")
-prediction = model(input)[0].numpy()[0]
+logits = model(input)[0].numpy()[0]
+prediction = logits / logits.sum()
 
 if st.button("Submit", type="primary"):
     label, probability = get_highest_toxicity_class(prediction)