Add to config.json the id2label and label2id keys.

README.md

@@ -34,7 +34,7 @@ with torch.no_grad():
   probs = model(**input).logits
 
 predicted_label = torch.argmax(probs, dim=1).item()
-print(predicted_label)  # 0 is neutral, 1 is gendered
+print(model.config.id2label[predicted_label])
 ```
 
 ## Citation

config.json

@@ -6,6 +6,14 @@
   "attention_probs_dropout_prob": 0.1,
   "bos_token_id": 5,
   "classifier_dropout": null,
+  "id2label": {
+    "0": "neutral",
+    "1": "gendered"
+  },
+  "label2id": {
+    "gendered": 1,
+    "neutral": 0
+  },
   "eos_token_id": 6,
   "hidden_act": "gelu",
   "hidden_dropout_prob": 0.1,