Update sroie_inference.py

sroie_inference.py

@@ -5,15 +5,14 @@ from PIL import Image, ImageDraw, ImageFont
 from transformers import LayoutLMv3TokenizerFast, LayoutLMv3Processor, LayoutLMv3ForTokenClassification
 from utils import OCR, unnormalize_box
 
-#  [B-COMPANY", "I-COMPANY", "B-DATE", "I-DATE", "B-ADDRESS", "I-ADDRESS", "B-TOTAL", "I-TOTAL", "O"]
-labels = ["COMPANY", "COMPANY", "DATE", "DATE", "ADDRESS", "ADDRESS", "TOTAL", "TOTAL", "O"]
-id2label = {v: k for v, k in enumerate(labels)}
-label2id = {k: v for v, k in enumerate(labels)}
 
 tokenizer = LayoutLMv3TokenizerFast.from_pretrained("mp-02/layoutlmv3-finetuned-sroie", apply_ocr=False)
 processor = LayoutLMv3Processor.from_pretrained("mp-02/layoutlmv3-finetuned-sroie", apply_ocr=False)
 model = LayoutLMv3ForTokenClassification.from_pretrained("mp-02/layoutlmv3-finetuned-sroie")
 
+id2label = model.config.id2label
+label2id = model.config.label2id
+
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model.to(device)