Update files/functions.py

files/functions.py

@@ -137,13 +137,21 @@ langdetect2Tesseract = {v:k for k,v in Tesseract2langdetect.items()}
 
 ## model / feature extractor / tokenizer
 
+# get device
 import torch
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-from transformers import LayoutLMv2ForTokenClassification # LayoutXLMTokenizerFast, 
+## model LiLT
+import transformers
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+tokenizer_lilt = AutoTokenizer.from_pretrained(model_id_lilt)
+model_lilt = AutoModelForTokenClassification.from_pretrained(model_id_lilt);
+model_lilt.to(device);
 
-model = LayoutLMv2ForTokenClassification.from_pretrained(model_id);
-model.to(device);
+## model LayoutXLM
+from transformers import LayoutLMv2ForTokenClassification # LayoutXLMTokenizerFast, 
+model_layoutxlm = LayoutLMv2ForTokenClassification.from_pretrained(model_id_layoutxlm);
+model_layoutxlm.to(device);
 
 # feature extractor
 from transformers import LayoutLMv2FeatureExtractor
@@ -151,12 +159,16 @@ feature_extractor = LayoutLMv2FeatureExtractor(apply_ocr=False)
 
 # tokenizer
 from transformers import AutoTokenizer
-tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
+tokenizer_layoutxlm = AutoTokenizer.from_pretrained(tokenizer_id_layoutxlm)
 
 # get labels
-id2label = model.config.id2label
-label2id = model.config.label2id
-num_labels = len(id2label)
+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