Update app.py
Browse files
app.py
CHANGED
@@ -30,8 +30,9 @@ def classify_text(text):
|
|
30 |
probabilities = torch.softmax(logits, dim=1).tolist()[0]
|
31 |
|
32 |
# Extract predicted labels
|
|
|
33 |
labels = ['Toxic', 'Severe Toxic', 'Obscene', 'Threat', 'Insult', 'Identity Hate']
|
34 |
-
predicted_labels = [labels[i] for i, prob in enumerate(probabilities) if prob >
|
35 |
|
36 |
return predicted_labels
|
37 |
|
@@ -65,7 +66,6 @@ def main():
|
|
65 |
if not df_results.empty:
|
66 |
st.subheader('Classification Results')
|
67 |
st.dataframe(df_results)
|
68 |
-
|
69 |
|
70 |
if __name__ == '__main__':
|
71 |
main()
|
|
|
30 |
probabilities = torch.softmax(logits, dim=1).tolist()[0]
|
31 |
|
32 |
# Extract predicted labels
|
33 |
+
threshold = 0.5
|
34 |
labels = ['Toxic', 'Severe Toxic', 'Obscene', 'Threat', 'Insult', 'Identity Hate']
|
35 |
+
predicted_labels = [labels[i] for i, prob in enumerate(probabilities) if prob > threshold]
|
36 |
|
37 |
return predicted_labels
|
38 |
|
|
|
66 |
if not df_results.empty:
|
67 |
st.subheader('Classification Results')
|
68 |
st.dataframe(df_results)
|
|
|
69 |
|
70 |
if __name__ == '__main__':
|
71 |
main()
|