dfucci g8a9 commited on
Commit
b3828fa
1 Parent(s): 9cf716e

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>

Files changed (2) hide show
  1. README.md +1 -1
  2. 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) # 0 is neutral, 1 is gendered
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,