1import re, json, torch, torch.nn as nn
2from dataclasses import dataclass
3from typing import List, Tuple, Optional
4from transformers import XLMRobertaModel, XLMRobertaConfig, PreTrainedModel, AutoTokenizer
5from transformers.modeling_outputs import ModelOutput
6from torchcrf import CRF
7
8# ── Label sets ───────────────────────────────────────────────────────────────
9ENTITY_TYPES = [
10 "STREET", "ROAD", "HOUSENUMBER", "POSTCODE", "CITY", "PROVINCE",
11 "BUILDING", "INTERSECTION", "PARCEL", "DISTRICT",
12 "GRAVE_LOCATION", "DOMAIN_ZONE_AREA",
13]
14BIO_LABELS = ["O"] + [f"{p}-{e}" for e in ENTITY_TYPES for p in ("B", "I")]
15LOC_BIO_LABELS = ["O", "B-LOCATION", "I-LOCATION"]
16label2id = {l: i for i, l in enumerate(BIO_LABELS)}
17id2label = {i: l for i, l in enumerate(BIO_LABELS)}
18loc_label2id = {l: i for i, l in enumerate(LOC_BIO_LABELS)}
19loc_id2label = {i: l for i, l in enumerate(LOC_BIO_LABELS)}
20MAX_LENGTH = 256
21
22# ── Model classes ─────────────────────────────────────────────────────────────
23class DualNERConfig(XLMRobertaConfig):
24 model_type = "dual_ner_xlm_roberta"
25 def __init__(self, num_component_labels=len(BIO_LABELS),
26 num_location_labels=len(LOC_BIO_LABELS), **kwargs):
27 super().__init__(**kwargs)
28 self.num_component_labels = num_component_labels
29 self.num_location_labels = num_location_labels
30
31@dataclass
32class DualNEROutput(ModelOutput):
33 loss: Optional[torch.FloatTensor] = None
34 component_logits: torch.FloatTensor = None
35 location_logits: torch.FloatTensor = None
36
37class DualHeadLocationNER(PreTrainedModel):
38 config_class = DualNERConfig
39 base_model_prefix = "roberta"
40
41 def __init__(self, config):
42 super().__init__(config)
43 self.roberta = XLMRobertaModel(config, add_pooling_layer=False)
44 self.dropout = nn.Dropout(config.hidden_dropout_prob)
45 self.component_head = nn.Sequential(
46 nn.Linear(config.hidden_size, 256), nn.GELU(), nn.Dropout(0.1),
47 nn.Linear(256, config.num_component_labels),
48 )
49 self.location_head = nn.Sequential(
50 nn.Linear(config.hidden_size, 256), nn.GELU(), nn.Dropout(0.1),
51 nn.Linear(256, config.num_location_labels),
52 )
53 self.component_crf = CRF(config.num_component_labels, batch_first=True)
54 self.location_crf = CRF(config.num_location_labels, batch_first=True)
55 self.post_init()
56
57 def forward(self, input_ids=None, attention_mask=None, **kwargs):
58 h = self.dropout(
59 self.roberta(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
60 )
61 return DualNEROutput(
62 component_logits=self.component_head(h),
63 location_logits=self.location_head(h),
64 )
65
66# ── Tokenizer helpers ─────────────────────────────────────────────────────────
67def tokenize_location(text: str) -> Tuple[List[str], List[Tuple[int, int]]]:
68 tokens, offsets = [], []
69 for m in re.finditer(r'[,;()\[\]{}]|[^\s,;()\[\]{}]+', text):
70 tokens.append(m.group())
71 offsets.append((m.start(), m.end()))
72 return tokens, offsets
73
74def classify_housenumber_type(hn: str) -> str:
75 if re.search(r'\d\s*[-–]\s*\d', hn): return "range"
76 if re.search(r'[,;]|\band\b|\ben\b', hn): return "sequence"
77 return "single"
78
79def extract_bio_spans(tokens, tag_ids, id2l, offsets):
80 spans, ent, start = [], None, 0
81 for i, tid in enumerate(tag_ids):
82 tag = id2l[tid]
83 if tag.startswith("B-"):
84 if ent: spans.append({"entity": ent, "start_tok": start, "end_tok": i - 1,
85 "char_start": offsets[start][0], "char_end": offsets[i-1][1]})
86 ent, start = tag[2:], i
87 elif not (tag.startswith("I-") and ent == tag[2:]):
88 if ent: spans.append({"entity": ent, "start_tok": start, "end_tok": i - 1,
89 "char_start": offsets[start][0], "char_end": offsets[i-1][1]})
90 ent = None
91 if ent: spans.append({"entity": ent, "start_tok": start, "end_tok": len(tokens)-1,
92 "char_start": offsets[start][0], "char_end": offsets[-1][1]})
93 return spans
94
95# ── Inference ─────────────────────────────────────────────────────────────────
96def predict_locations(text, model, tokenizer, device="cpu"):
97 tokens, offsets = tokenize_location(text)
98 if not tokens:
99 return {"original": text, "locations": []}
100 enc = tokenizer(tokens, is_split_into_words=True, return_tensors="pt",
101 truncation=True, max_length=MAX_LENGTH)
102 word_ids = enc.word_ids()
103 with torch.no_grad():
104 out = model(**{k: v.to(device) for k, v in enc.items()})
105 mask = enc["attention_mask"].bool().to(device)
106 cpreds = model.component_crf.decode(out.component_logits, mask=mask)[0]
107 lpreds = model.location_crf.decode(out.location_logits, mask=mask)[0]
108 wcomp, wloc, prev = [], [], None
109 for idx, wid in enumerate(word_ids):
110 if wid is None: continue
111 if wid != prev:
112 wcomp.append(cpreds[idx]); wloc.append(lpreds[idx])
113 prev = wid
114 loc_spans = extract_bio_spans(tokens, wloc, loc_id2label, offsets)
115 comp_spans = extract_bio_spans(tokens, wcomp, id2label, offsets)
116 locations, assigned = [], set()
117 for ls in loc_spans:
118 loc = {"location": text[ls["char_start"]:ls["char_end"]]}
119 for ci, cs in enumerate(comp_spans):
120 if cs["start_tok"] >= ls["start_tok"] and cs["end_tok"] <= ls["end_tok"]:
121 loc[cs["entity"].lower()] = text[cs["char_start"]:cs["char_end"]]
122 assigned.add(ci)
123 if "housenumber" in loc:
124 loc["housenumber_type"] = classify_housenumber_type(loc["housenumber"])
125 locations.append(loc)
126 for ci, cs in enumerate(comp_spans):
127 if ci not in assigned:
128 loc = {"location": text[cs["char_start"]:cs["char_end"]],
129 cs["entity"].lower(): text[cs["char_start"]:cs["char_end"]]}
130 if "housenumber" in loc:
131 loc["housenumber_type"] = classify_housenumber_type(loc["housenumber"])
132 locations.append(loc)
133 return {"original": text, "locations": locations}
134
135# ── Load & run ────────────────────────────────────────────────────────────────
136MODEL_REPO = "svercoutere/abb-dual-location-component-ner"
137device = "cuda" if torch.cuda.is_available() else "cpu"
138
139tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
140config = DualNERConfig.from_pretrained(MODEL_REPO,
141 num_component_labels=len(BIO_LABELS),
142 num_location_labels=len(LOC_BIO_LABELS))
143model = DualHeadLocationNER.from_pretrained(MODEL_REPO, config=config)
144model.to(device).eval()
145
146texts = [
147 "Scaldisstraat 23-25, 2000 Antwerpen",
148 "Cafe den Draak, Lovegemlaan 7, 9000 Gent",
149 "Heikeesstraat 2, 9240 Zele and Dorpstraat 7, 8040 Mariakerke",
150 "begraafplaats Schoonselhof, perk 27, rij 3",
151 "politiezone Antwerpen",
152]
153for text in texts:
154 result = predict_locations(text, model, tokenizer, device)
155 print(f"\nInput : {text}")
156 for loc in result["locations"]:
157 parts = {k: v for k, v in loc.items() if k != "location"}
158 print(f" LOC : {loc['location']}")
159 print(f" {json.dumps(parts, ensure_ascii=False)}")
