pierreguillou commited on
Commit
e19b360
1 Parent(s): d3e0af2

Update files/functions.py

Browse files
Files changed (1) hide show
  1. files/functions.py +5 -0
files/functions.py CHANGED
@@ -157,6 +157,11 @@ from transformers import AutoTokenizer
157
  tokenizer_id = "xlm-roberta-base"
158
  tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
159
 
 
 
 
 
 
160
  ## General
161
 
162
  # get text and bounding boxes from an image
 
157
  tokenizer_id = "xlm-roberta-base"
158
  tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
159
 
160
+ # get labels
161
+ id2label = model.config.id2label
162
+ label2id = model.config.label2id
163
+ num_labels = len(id2label)
164
+
165
  ## General
166
 
167
  # get text and bounding boxes from an image