Update README.md

README.md

@@ -1,36 +1,36 @@
----
-license: apache-2.0
-language:
-- en
-- fr
-- de
-- it
-- es
-- pt
-- pl
-- zh
-- nl
-base_model:
-- FacebookAI/xlm-roberta-base
-datasets:
-- PleIAs/ToxicCommons
-- yangezheng/SWSR-SexComment
-tags:
-- toxicity
-- data
----
+---
+license: apache-2.0
+language:
+- en
+- fr
+- de
+- it
+- es
+- pt
+- pl
+- zh
+- nl
+base_model:
+- FacebookAI/xlm-roberta-base
+datasets:
+- PleIAs/ToxicCommons
+- yangezheng/SWSR-SexComment
+tags:
+- toxicity
+- data
+---
 # Multilingual Toxicity Classifiers used in Apertus Pretraining
 #### Author: Olivia Simin Fan (@Olivia-umich)
 
 Language specific toxicity classifiers in English, French, German, Italian, Spanish, Portuguese, Polish, Chinese and Dutch, trained on [PleIAs/ToxicCommons](https://github.com/Pleias/toxic-commons) and [SWSR-SexComments](https://arxiv.org/pdf/2108.03070) datasets.
-
+We subsample non-toxic examples to create balanced 50%-50% training sets for each language. We separate 10% from the balanced dataset to build a balanced held-out validation set.
 
 ## Model Description
 Our toxicity classifier employs a two-stage approach: we first extract the multilingual document embeddings using [*XLM-RoBERTa*](https://huggingface.co/FacebookAI/xlm-roberta-base), 
 then train a language-specific 2-layer MLP for binary toxicity classification on top of these embeddings for 6 epochs. 
 The classifier checkpoints with the best accuracy on the held-out validation set are further employed to annotate the toxicity scores on FineWeb-2 and FineWeb.
 
-The validation accuracies on the held-out test set is provided as below:
+The validation accuracies on the balanced held-out test set are provided as below:
 | Language | Accuracy |
 |----------|----------|
 | English (en) | 80.13% |
@@ -110,7 +110,8 @@ class RobertaClassifier(nn.Module):
 
 
 ```python
-MODEL_PATH = f"{MODEL_DIR}/english.pth"
+LANGUAGE = "english" # choose from ["english", "chinese", "french", "german", "italian", "spanish", "portuguese", "polish", "dutch"]
+MODEL_PATH = f"{MODEL_DIR}/{LANGUAGE}.pth"
 DEVICE = "cpu"
 
 model = RobertaClassifier(device=DEVICE, num_classes=2)