Update README.md

README.md

@@ -11,40 +11,32 @@ The model is translated from [PaddlePaddle/ernie-layoutx-base-uncased](https://h
 ```python
 import torch
 from PIL import Image
-import numpy as np
 import torch.nn.functional as F
-from networks.model_util import ernie_qa_processing
-from networks import ErnieLayoutConfig, ErnieLayoutForQuestionAnswering, ErnieLayoutImageProcessor, \
-    ERNIELayoutProcessor, ErnieLayoutTokenizerFast
+from networks import ErnieLayoutConfig, ErnieLayoutForQuestionAnswering, \
+    ErnieLayoutProcessor, ErnieLayoutTokenizerFast
+from transformers.models.layoutlmv3 import LayoutLMv3ImageProcessor
 
 pretrain_torch_model_or_path = "Norm/ERNIE-Layout-Pytorch"
-doc_imag_path = "/path/to/dummy_input.jpeg"
+doc_imag_path = "./dummy_input.jpeg"
 
-device = torch.device("cuda:0")
-
-# Dummy Input
-context = ['This is an example document', 'All ocr boxes are inserted into this list']
-layout = [[381, 91, 505, 115], [738, 96, 804, 122]]  # all boxes are resized between 0 - 1000
+context = ['This is an example sequence', 'All ocr boxes are inserted into this list']
+layout = [[381, 91, 505, 115], [738, 96, 804, 122]]  # make sure  all boxes are normalized between 0 - 1000
 pil_image = Image.open(doc_imag_path).convert("RGB")
 
 # initialize tokenizer
 tokenizer = ErnieLayoutTokenizerFast.from_pretrained(pretrained_model_name_or_path=pretrain_torch_model_or_path)
 
 # initialize feature extractor
-feature_extractor = ErnieLayoutImageProcessor(apply_ocr=False)
-processor = ERNIELayoutProcessor(image_processor=feature_extractor, tokenizer=tokenizer)
+feature_extractor = LayoutLMv3ImageProcessor(apply_ocr=False)
+processor = ErnieLayoutProcessor(image_processor=feature_extractor, tokenizer=tokenizer)
 
 # Tokenize context & questions
-context_encodings = processor(pil_image, context)
 question = "what is it?"
-tokenized_res = ernie_qa_processing(tokenizer, question, layout, context_encodings)
-tokenized_res['input_ids'] = torch.tensor([tokenized_res['input_ids']]).to(device)
-tokenized_res['bbox'] = torch.tensor([tokenized_res['bbox']]).to(device)
-tokenized_res['pixel_values'] = torch.tensor(np.array(context_encodings.data['pixel_values'])).to(device)
+encoding = processor(pil_image, question, context, boxes=layout, return_tensors="pt")
 
 # dummy answer start && end index
-tokenized_res['start_positions'] = torch.tensor([6]).to(device)
-tokenized_res['end_positions'] = torch.tensor([12]).to(device)
+start_positions = torch.tensor([6])
+end_positions = torch.tensor([12])
 
 # initialize config
 config = ErnieLayoutConfig.from_pretrained(pretrained_model_name_or_path=pretrain_torch_model_or_path)
@@ -55,15 +47,15 @@ model = ErnieLayoutForQuestionAnswering.from_pretrained(
     pretrained_model_name_or_path=pretrain_torch_model_or_path,
     config=config,
 )
-model.to(device)
 
-output = model(**tokenized_res)
+output = model(**encoding, start_positions=start_positions, end_positions=end_positions)
 
 # decode output
 start_max = torch.argmax(F.softmax(output.start_logits, dim=-1))
 end_max = torch.argmax(F.softmax(output.end_logits, dim=-1)) + 1  # add one ##because of python list indexing
-answer = tokenizer.decode(tokenized_res["input_ids"][0][start_max: end_max])
+answer = tokenizer.decode(encoding.input_ids[0][start_max: end_max])
 print(answer)
 
 
+
 ```