Add to config.json the id2label and label2id keys. (#3)
Browse files- Update config.json (2bb477fb4f29e057c4d566b76f83506d6302ad54)
- update usage example (28dcf06e6449afac6f0ff4b981ea20d1f1703a5b)
Co-authored-by: Giuseppe Attanasio <g8a9@users.noreply.huggingface.co>
- README.md +1 -1
- config.json +8 -0
README.md
CHANGED
@@ -34,7 +34,7 @@ with torch.no_grad():
|
|
34 |
probs = model(**input).logits
|
35 |
|
36 |
predicted_label = torch.argmax(probs, dim=1).item()
|
37 |
-
print(predicted_label)
|
38 |
```
|
39 |
|
40 |
## Citation
|
|
|
34 |
probs = model(**input).logits
|
35 |
|
36 |
predicted_label = torch.argmax(probs, dim=1).item()
|
37 |
+
print(model.config.id2label[predicted_label])
|
38 |
```
|
39 |
|
40 |
## Citation
|
config.json
CHANGED
@@ -6,6 +6,14 @@
|
|
6 |
"attention_probs_dropout_prob": 0.1,
|
7 |
"bos_token_id": 5,
|
8 |
"classifier_dropout": null,
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
9 |
"eos_token_id": 6,
|
10 |
"hidden_act": "gelu",
|
11 |
"hidden_dropout_prob": 0.1,
|
|
|
6 |
"attention_probs_dropout_prob": 0.1,
|
7 |
"bos_token_id": 5,
|
8 |
"classifier_dropout": null,
|
9 |
+
"id2label": {
|
10 |
+
"0": "neutral",
|
11 |
+
"1": "gendered"
|
12 |
+
},
|
13 |
+
"label2id": {
|
14 |
+
"gendered": 1,
|
15 |
+
"neutral": 0
|
16 |
+
},
|
17 |
"eos_token_id": 6,
|
18 |
"hidden_act": "gelu",
|
19 |
"hidden_dropout_prob": 0.1,
|