]> git.openstreetmap.org Git - nominatim.git/blob - nominatim/api/search/icu_tokenizer.py
Merge remote-tracking branch 'upstream/master'
[nominatim.git] / nominatim / api / search / icu_tokenizer.py
1 # SPDX-License-Identifier: GPL-3.0-or-later
2 #
3 # This file is part of Nominatim. (https://nominatim.org)
4 #
5 # Copyright (C) 2023 by the Nominatim developer community.
6 # For a full list of authors see the git log.
7 """
8 Implementation of query analysis for the ICU tokenizer.
9 """
10 from typing import Tuple, Dict, List, Optional, NamedTuple, Iterator, Any, cast
11 from collections import defaultdict
12 import dataclasses
13 import difflib
14
15 from icu import Transliterator
16
17 import sqlalchemy as sa
18
19 from nominatim.typing import SaRow
20 from nominatim.api.connection import SearchConnection
21 from nominatim.api.logging import log
22 from nominatim.api.search import query as qmod
23 from nominatim.api.search.query_analyzer_factory import AbstractQueryAnalyzer
24 from nominatim.db.sqlalchemy_types import Json
25
26
27 DB_TO_TOKEN_TYPE = {
28     'W': qmod.TokenType.WORD,
29     'w': qmod.TokenType.PARTIAL,
30     'H': qmod.TokenType.HOUSENUMBER,
31     'P': qmod.TokenType.POSTCODE,
32     'C': qmod.TokenType.COUNTRY
33 }
34
35
36 class QueryPart(NamedTuple):
37     """ Normalized and transliterated form of a single term in the query.
38         When the term came out of a split during the transliteration,
39         the normalized string is the full word before transliteration.
40         The word number keeps track of the word before transliteration
41         and can be used to identify partial transliterated terms.
42     """
43     token: str
44     normalized: str
45     word_number: int
46
47
48 QueryParts = List[QueryPart]
49 WordDict = Dict[str, List[qmod.TokenRange]]
50
51 def yield_words(terms: List[QueryPart], start: int) -> Iterator[Tuple[str, qmod.TokenRange]]:
52     """ Return all combinations of words in the terms list after the
53         given position.
54     """
55     total = len(terms)
56     for first in range(start, total):
57         word = terms[first].token
58         yield word, qmod.TokenRange(first, first + 1)
59         for last in range(first + 1, min(first + 20, total)):
60             word = ' '.join((word, terms[last].token))
61             yield word, qmod.TokenRange(first, last + 1)
62
63
64 @dataclasses.dataclass
65 class ICUToken(qmod.Token):
66     """ Specialised token for ICU tokenizer.
67     """
68     word_token: str
69     info: Optional[Dict[str, Any]]
70
71     def get_category(self) -> Tuple[str, str]:
72         assert self.info
73         return self.info.get('class', ''), self.info.get('type', '')
74
75
76     def rematch(self, norm: str) -> None:
77         """ Check how well the token matches the given normalized string
78             and add a penalty, if necessary.
79         """
80         if not self.lookup_word:
81             return
82
83         seq = difflib.SequenceMatcher(a=self.lookup_word, b=norm)
84         distance = 0
85         for tag, afrom, ato, bfrom, bto in seq.get_opcodes():
86             if tag in ('delete', 'insert') and (afrom == 0 or ato == len(self.lookup_word)):
87                 distance += 1
88             elif tag == 'replace':
89                 distance += max((ato-afrom), (bto-bfrom))
90             elif tag != 'equal':
91                 distance += abs((ato-afrom) - (bto-bfrom))
92         self.penalty += (distance/len(self.lookup_word))
93
94
95     @staticmethod
96     def from_db_row(row: SaRow) -> 'ICUToken':
97         """ Create a ICUToken from the row of the word table.
98         """
99         count = 1 if row.info is None else row.info.get('count', 1)
100         addr_count = 1 if row.info is None else row.info.get('addr_count', 1)
101
102         penalty = 0.0
103         if row.type == 'w':
104             penalty = 0.3
105         elif row.type == 'W':
106             if len(row.word_token) == 1 and row.word_token == row.word:
107                 penalty = 0.2 if row.word.isdigit() else 0.3
108         elif row.type == 'H':
109             penalty = sum(0.1 for c in row.word_token if c != ' ' and not c.isdigit())
110             if all(not c.isdigit() for c in row.word_token):
111                 penalty += 0.2 * (len(row.word_token) - 1)
112         elif row.type == 'C':
113             if len(row.word_token) == 1:
114                 penalty = 0.3
115
116         if row.info is None:
117             lookup_word = row.word
118         else:
119             lookup_word = row.info.get('lookup', row.word)
120         if lookup_word:
121             lookup_word = lookup_word.split('@', 1)[0]
122         else:
123             lookup_word = row.word_token
124
125         return ICUToken(penalty=penalty, token=row.word_id, count=max(1, count),
126                         lookup_word=lookup_word, is_indexed=True,
127                         word_token=row.word_token, info=row.info,
128                         addr_count=max(1, addr_count))
129
130
131
132 class ICUQueryAnalyzer(AbstractQueryAnalyzer):
133     """ Converter for query strings into a tokenized query
134         using the tokens created by a ICU tokenizer.
135     """
136
137     def __init__(self, conn: SearchConnection) -> None:
138         self.conn = conn
139
140
141     async def setup(self) -> None:
142         """ Set up static data structures needed for the analysis.
143         """
144         async def _make_normalizer() -> Any:
145             rules = await self.conn.get_property('tokenizer_import_normalisation')
146             return Transliterator.createFromRules("normalization", rules)
147
148         self.normalizer = await self.conn.get_cached_value('ICUTOK', 'normalizer',
149                                                            _make_normalizer)
150
151         async def _make_transliterator() -> Any:
152             rules = await self.conn.get_property('tokenizer_import_transliteration')
153             return Transliterator.createFromRules("transliteration", rules)
154
155         self.transliterator = await self.conn.get_cached_value('ICUTOK', 'transliterator',
156                                                                _make_transliterator)
157
158         if 'word' not in self.conn.t.meta.tables:
159             sa.Table('word', self.conn.t.meta,
160                      sa.Column('word_id', sa.Integer),
161                      sa.Column('word_token', sa.Text, nullable=False),
162                      sa.Column('type', sa.Text, nullable=False),
163                      sa.Column('word', sa.Text),
164                      sa.Column('info', Json))
165
166
167     async def analyze_query(self, phrases: List[qmod.Phrase]) -> qmod.QueryStruct:
168         """ Analyze the given list of phrases and return the
169             tokenized query.
170         """
171         log().section('Analyze query (using ICU tokenizer)')
172         normalized = list(filter(lambda p: p.text,
173                                  (qmod.Phrase(p.ptype, self.normalize_text(p.text))
174                                   for p in phrases)))
175         query = qmod.QueryStruct(normalized)
176         log().var_dump('Normalized query', query.source)
177         if not query.source:
178             return query
179
180         parts, words = self.split_query(query)
181         log().var_dump('Transliterated query', lambda: _dump_transliterated(query, parts))
182
183         for row in await self.lookup_in_db(list(words.keys())):
184             for trange in words[row.word_token]:
185                 token = ICUToken.from_db_row(row)
186                 if row.type == 'S':
187                     if row.info['op'] in ('in', 'near'):
188                         if trange.start == 0:
189                             query.add_token(trange, qmod.TokenType.NEAR_ITEM, token)
190                     else:
191                         if trange.start == 0 and trange.end == query.num_token_slots():
192                             query.add_token(trange, qmod.TokenType.NEAR_ITEM, token)
193                         else:
194                             query.add_token(trange, qmod.TokenType.QUALIFIER, token)
195                 else:
196                     query.add_token(trange, DB_TO_TOKEN_TYPE[row.type], token)
197
198         self.add_extra_tokens(query, parts)
199         self.rerank_tokens(query, parts)
200
201         log().table_dump('Word tokens', _dump_word_tokens(query))
202
203         return query
204
205
206     def normalize_text(self, text: str) -> str:
207         """ Bring the given text into a normalized form. That is the
208             standardized form search will work with. All information removed
209             at this stage is inevitably lost.
210         """
211         norm = cast(str, self.normalizer.transliterate(text))
212         numspaces = norm.count(' ')
213         if numspaces > 4 and len(norm) <= (numspaces + 1) * 3:
214             return ''
215
216         return norm
217
218
219     def split_query(self, query: qmod.QueryStruct) -> Tuple[QueryParts, WordDict]:
220         """ Transliterate the phrases and split them into tokens.
221
222             Returns the list of transliterated tokens together with their
223             normalized form and a dictionary of words for lookup together
224             with their position.
225         """
226         parts: QueryParts = []
227         phrase_start = 0
228         words = defaultdict(list)
229         wordnr = 0
230         for phrase in query.source:
231             query.nodes[-1].ptype = phrase.ptype
232             for word in phrase.text.split(' '):
233                 trans = self.transliterator.transliterate(word)
234                 if trans:
235                     for term in trans.split(' '):
236                         if term:
237                             parts.append(QueryPart(term, word, wordnr))
238                             query.add_node(qmod.BreakType.TOKEN, phrase.ptype)
239                     query.nodes[-1].btype = qmod.BreakType.WORD
240                 wordnr += 1
241             query.nodes[-1].btype = qmod.BreakType.PHRASE
242
243             for word, wrange in yield_words(parts, phrase_start):
244                 words[word].append(wrange)
245
246             phrase_start = len(parts)
247         query.nodes[-1].btype = qmod.BreakType.END
248
249         return parts, words
250
251
252     async def lookup_in_db(self, words: List[str]) -> 'sa.Result[Any]':
253         """ Return the token information from the database for the
254             given word tokens.
255         """
256         t = self.conn.t.meta.tables['word']
257         return await self.conn.execute(t.select().where(t.c.word_token.in_(words)))
258
259
260     def add_extra_tokens(self, query: qmod.QueryStruct, parts: QueryParts) -> None:
261         """ Add tokens to query that are not saved in the database.
262         """
263         for part, node, i in zip(parts, query.nodes, range(1000)):
264             if len(part.token) <= 4 and part[0].isdigit()\
265                and not node.has_tokens(i+1, qmod.TokenType.HOUSENUMBER):
266                 query.add_token(qmod.TokenRange(i, i+1), qmod.TokenType.HOUSENUMBER,
267                                 ICUToken(0.5, 0, 1, 1, part.token, True, part.token, None))
268
269
270     def rerank_tokens(self, query: qmod.QueryStruct, parts: QueryParts) -> None:
271         """ Add penalties to tokens that depend on presence of other token.
272         """
273         for i, node, tlist in query.iter_token_lists():
274             if tlist.ttype == qmod.TokenType.POSTCODE:
275                 for repl in node.starting:
276                     if repl.end == tlist.end and repl.ttype != qmod.TokenType.POSTCODE \
277                        and (repl.ttype != qmod.TokenType.HOUSENUMBER
278                             or len(tlist.tokens[0].lookup_word) > 4):
279                         repl.add_penalty(0.39)
280             elif tlist.ttype == qmod.TokenType.HOUSENUMBER \
281                  and len(tlist.tokens[0].lookup_word) <= 3:
282                 if any(c.isdigit() for c in tlist.tokens[0].lookup_word):
283                     for repl in node.starting:
284                         if repl.end == tlist.end and repl.ttype != qmod.TokenType.HOUSENUMBER:
285                             repl.add_penalty(0.5 - tlist.tokens[0].penalty)
286             elif tlist.ttype not in (qmod.TokenType.COUNTRY, qmod.TokenType.PARTIAL):
287                 norm = parts[i].normalized
288                 for j in range(i + 1, tlist.end):
289                     if parts[j - 1].word_number != parts[j].word_number:
290                         norm += '  ' + parts[j].normalized
291                 for token in tlist.tokens:
292                     cast(ICUToken, token).rematch(norm)
293
294
295 def _dump_transliterated(query: qmod.QueryStruct, parts: QueryParts) -> str:
296     out = query.nodes[0].btype.value
297     for node, part in zip(query.nodes[1:], parts):
298         out += part.token + node.btype.value
299     return out
300
301
302 def _dump_word_tokens(query: qmod.QueryStruct) -> Iterator[List[Any]]:
303     yield ['type', 'token', 'word_token', 'lookup_word', 'penalty', 'count', 'info']
304     for node in query.nodes:
305         for tlist in node.starting:
306             for token in tlist.tokens:
307                 t = cast(ICUToken, token)
308                 yield [tlist.ttype.name, t.token, t.word_token or '',
309                        t.lookup_word or '', t.penalty, t.count, t.info]
310
311
312 async def create_query_analyzer(conn: SearchConnection) -> AbstractQueryAnalyzer:
313     """ Create and set up a new query analyzer for a database based
314         on the ICU tokenizer.
315     """
316     out = ICUQueryAnalyzer(conn)
317     await out.setup()
318
319     return out