Views
No views yet
1from transformers import AutoTokenizer, AutoModelForQuestionAnswering
2from datasets import load_dataset
3
4model_checkpoint = "TusharGoel/LiLT-Document-QA"
5tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, add_prefix_space=True)
6model_predict = AutoModelForQuestionAnswering.from_pretrained(model_checkpoint)
7
8model_predict.eval()
9dataset = load_dataset("nielsr/funsd", split="train")
10example = dataset[0]
11print(example)
12
13question = "What is the Licensee Number?"
14print(question)
15
16words = example["words"]
17boxes = example["bboxes"]
18
19encoding = tokenizer(question, words, boxes = boxes, return_token_type_ids=True, return_tensors="pt")
20
21word_ids = encoding.word_ids(0)
22outputs = model_predict(**encoding)
23
24loss = outputs.loss
25start_scores = outputs.start_logits
26end_scores = outputs.end_logits
27
28start, end = word_ids[start_scores.argmax(-1).item()], word_ids[end_scores.argmax(-1).item()]
29# print(start, end)
30print(" ".join(words[start : end + 1]))