Upload toxic.py

toxic.py

@@ -0,0 +1,29 @@
+import streamlit as st
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+model_checkpoint = 'cointegrated/rubert-tiny-toxicity'
+tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
+model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint)
+if torch.cuda.is_available():
+    model.cuda()
+
+def text2toxicity(text, aggregate=True):
+    """ Calculate toxicity of a text (if aggregate=True) or a vector of toxicity aspects (if aggregate=False)"""
+    with torch.no_grad():
+        inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True).to(model.device)
+        proba = torch.sigmoid(model(**inputs).logits).cpu().numpy()
+    if isinstance(text, str):
+        proba = proba[0]
+    if aggregate:
+        return 1 - proba.T[0] * (1 - proba.T[-1])
+    return proba
+
+st.title("Toxicity Detector")
+
+user_input = st.text_area("Enter text to check for toxicity:", "Капец ты гнида")
+if st.button("Analyze"):
+    toxicity_score = text2toxicity(user_input, True)
+    st.write(f"Toxicity Score: {toxicity_score:.4f}")
+    if toxicity_score > 0.5:
+        st.write("Warning: The text seems to be toxic!")