Update app.py

app.py

@@ -30,8 +30,9 @@ def classify_text(text):
         probabilities = torch.softmax(logits, dim=1).tolist()[0]
     
     # Extract predicted labels
+    threshold = 0.5
     labels = ['Toxic', 'Severe Toxic', 'Obscene', 'Threat', 'Insult', 'Identity Hate']
-    predicted_labels = [labels[i] for i, prob in enumerate(probabilities) if prob > 0.5]
+    predicted_labels = [labels[i] for i, prob in enumerate(probabilities) if prob > threshold]
     
     return predicted_labels
 
@@ -65,7 +66,6 @@ def main():
     if not df_results.empty:
         st.subheader('Classification Results')
         st.dataframe(df_results)
-    
 
 if __name__ == '__main__':
     main()