Update app.py

app.py

@@ -9,27 +9,21 @@ from transformers import AutoProcessor
 from datasets import Features, Sequence, ClassLabel, Value, Array2D, Array3D
 import torch
 from datasets import load_metric
-from transformers import LayoutLMv3ForTokenClassification
+from transformers import LayoutLMTokenizer
 from transformers.data.data_collator import default_data_collator
 
 
-from transformers import AutoModelForTokenClassification
+from transformers import LayoutLMForTokenClassification
 from datasets import load_dataset
 from PIL import Image, ImageDraw, ImageFont
 
 
-processor = AutoProcessor.from_pretrained("microsoft/layoutlmv3-base", apply_ocr=True)
-model = AutoModelForTokenClassification.from_pretrained("oussama/Layoutlm_Form_information_extraction")
+tokenizer = LayoutLMTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
+model = LayoutLMForTokenClassification.from_pretrained("microsoft/layoutlm-base-uncased", num_labels=13)
+
 
 
 
-# load image example
-dataset = load_dataset("darentang/generated", split="test")
-Image.open(dataset[2]["image_path"]).convert("RGB").save("example1.png")
-Image.open(dataset[1]["image_path"]).convert("RGB").save("example2.png")
-Image.open(dataset[0]["image_path"]).convert("RGB").save("example3.png")
-# define id2label, label2color
-labels = dataset.features['ner_tags'].feature.names
 id2label = {0: 'O', 1: 'B-HEADER', 2: 'I-HEADER', 3: 'B-QUESTION', 4: 'I-QUESTION', 5: 'B-ANSWER', 6: 'I-ANSWER'}
 label2color = {'question':'blue', 'answer':'green', 'header':'orange', 'other':'violet'}
 
@@ -54,7 +48,7 @@ def process_image(image):
     width, height = image.size
 
     # encode
-    encoding = processor(image, truncation=True, return_offsets_mapping=True, return_tensors="pt")
+    encoding = tokenizer(image, truncation=True, return_offsets_mapping=True, return_tensors="pt")
     offset_mapping = encoding.pop('offset_mapping')
 
     # forward pass