Update files/functions.py

files/functions.py

@@ -172,6 +172,15 @@ feature_extractor = LayoutLMv2FeatureExtractor(apply_ocr=False)
 from transformers import AutoTokenizer
 tokenizer_layoutxlm = AutoTokenizer.from_pretrained(tokenizer_id_layoutxlm)
 
+# get labels
+id2label_lilt = model_lilt.config.id2label
+label2id_lilt = model_lilt.config.label2id
+num_labels_lilt = len(id2label_lilt)
+
+id2label_layoutxlm = model_layoutxlm.config.id2label
+label2id_layoutxlm = model_layoutxlm.config.label2id
+num_labels_layoutxlm = len(id2label_layoutxlm)
+
 ## General
 
 # get text and bounding boxes from an image