Views
No views yet
1
2id2label = {0: "none", 1:"title", 2:"content", 3:"author", 4: "date", 5: "header", 6: "footer", 7: "rail", 8: "advertisement", 9: "navigation"}
3
4def eval(url):
5 current_dir = os.path.dirname(os.path.abspath(__file__))
6
7 model_folder = os.path.join(current_dir, 'models') # models folder is in the repository root
8 model_name = 'OxMarkupLM.pt'
9
10 processor = MarkupLMProcessor.from_pretrained("microsoft/markuplm-base")
11 processor.parse_html = False
12
13 model_path = os.path.join(model_folder, model_name)
14
15 model = MarkupLMForTokenClassification.from_pretrained(
16 model_path, id2label=labels.id2label, label2id=labels.label2id
17 )
18
19 html = utils.clean_html(utils.get_html_content(url))
20 data = [utils.extract_nodes_and_feautures(html)]
21 example = utils.split_sliding_data(data, 10, 0)
22
23 title, author, date, content = [], [], [], []
24 for splited in example:
25 nodes, xpaths = splited['nodes'], splited['xpaths']
26 encoding = processor(
27 nodes=nodes, xpaths=xpaths, return_offsets_mapping=True,
28 padding="max_length", truncation=True, max_length=512, return_tensors="pt"
29 )
30 offset_mapping = encoding.pop("offset_mapping")
31 with torch.no_grad():
32 logits = model(**encoding).logits
33
34 predictions = logits.argmax(-1)
35 processed_words = []
36
37 for pred_id, word_id, offset in zip(predictions[0].tolist(), encoding.word_ids(0), offset_mapping[0].tolist()):
38 if word_id is not None and offset[0] == 0:
39 if pred_id == 1:
40 title.append(nodes[word_id])
41 elif pred_id == 2 and word_id not in processed_words:
42 processed_words.append(word_id)
43 content.append(nodes[word_id])
44 elif pred_id == 3:
45 author.append(nodes[word_id])
46 elif pred_id == 4:
47 date.append(nodes[word_id])
48
49 title = rank_titles(title, '\n'.join(content))
50 return {
51 "model_name": model_name,
52 "url": url,
53 "title": title,
54 "author": author,
55 "date": date,
56 "content": content,
57 }