NielsRogge
commented
1 year ago Hi,
To perform inference with LiLT, you don't need a processor, as the model only gets text and corresponding bounding boxes as input. We only need a tokenizer.
Inference can be performed as follows:
from transformers import AutoTokenizer, AutoModelForTokenClassification
from datasets import load_dataset
tokenizer = AutoTokenizer.from_pretrained("nielsr/lilt-xlm-roberta-base")
model = AutoModelForTokenClassification.from_pretrained("nielsr/lilt-xlm-roberta-base")
# load dataset
dataset = load_dataset("nielsr/funsd-iob-original", split="train")
# load an image on which we'd like to predict,
# along with the words + coordinates recognized by an OCR engine
example = self.dataset[0]
image = example["image"]
words = example["words"]
boxes = example["original_bboxes"]
# prepare for the model
width, height = image.size
bbox = []
for word, box, label in zip(words, boxes, ner_tags):
box = normalize_bbox(box, width, height)
n_word_tokens = len(tokenizer.tokenize(word))
bbox.extend([box] * n_word_tokens)
# add special tokens
cls_box = sep_box = [0, 0, 0, 0]
bbox = [cls_box] + bbox + [sep_box]
encoding = self.tokenizer(" ".join(words), truncation=True, max_length=512)
sequence_length = len(encoding.input_ids)
# truncate boxes based on length of input ids
bbox = bbox[:sequence_length]
encoding["bbox"] = bboxNext you can do a forward pass and decode the predictions:
import torch
# forward pass
with torch.no_grad():
outputs = model(**encoding)
predicted_class_indices = outputs.logits.argmax(-1)[0].tolist()
# turn into actual class names
predicted_classes = [model.config.id2label[label] for label in predicted_class_indices]Note: I'm loading a model with a randomly initialized classification head above, hence predictions will be random. You need to load a model which has a fine-tuned classification head.
See also my LayoutLMv1 notebooks regarding turning the token-level predictions into word-level predictions and visualization.
piegu
eschaffn
mellahysf
philmas
Hi @NielsRogge,
in your notebook Fine_tune_LiLT_on_a_custom_dataset%2C_in_any_language.ipynb, you do not write the code for inference from a document image.
Can you give it please? I did try to copy the code from "Document AI: LiLT a better language agnostic LayoutLM model" but the fact to use an
xlm-robertaand not arobertaseems to complicate the processor instance.