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 copy import copy
  12 from collections import defaultdict
  13 import dataclasses
  14 import difflib
  15
  16 from icu import Transliterator
  17
  18 import sqlalchemy as sa
  19
  20 from nominatim.typing import SaRow
  21 from nominatim.api.connection import SearchConnection
  22 from nominatim.api.logging import log
  23 from nominatim.api.search import query as qmod
  24 from nominatim.api.search.query_analyzer_factory import AbstractQueryAnalyzer
  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
 101         penalty = 0.0
 102         if row.type == 'w':
 103             penalty = 0.3
 104         elif row.type == 'W':
 105             if len(row.word_token) == 1 and row.word_token == row.word:
 106                 penalty = 0.2 if row.word.isdigit() else 0.3
 107         elif row.type == 'H':
 108             penalty = sum(0.1 for c in row.word_token if c != ' ' and not c.isdigit())
 109             if all(not c.isdigit() for c in row.word_token):
 110                 penalty += 0.2 * (len(row.word_token) - 1)
 111         elif row.type == 'C':
 112             if len(row.word_token) == 1:
 113                 penalty = 0.3
 114
 115         if row.info is None:
 116             lookup_word = row.word
 117         else:
 118             lookup_word = row.info.get('lookup', row.word)
 119         if lookup_word:
 120             lookup_word = lookup_word.split('@', 1)[0]
 121         else:
 122             lookup_word = row.word_token
 123
 124         return ICUToken(penalty=penalty, token=row.word_id, count=count,
 125                         lookup_word=lookup_word, is_indexed=True,
 126                         word_token=row.word_token, info=row.info)
 127
 128
 129
 130 class ICUQueryAnalyzer(AbstractQueryAnalyzer):
 131     """ Converter for query strings into a tokenized query
 132         using the tokens created by a ICU tokenizer.
 133     """
 134
 135     def __init__(self, conn: SearchConnection) -> None:
 136         self.conn = conn
 137
 138
 139     async def setup(self) -> None:
 140         """ Set up static data structures needed for the analysis.
 141         """
 142         async def _make_normalizer() -> Any:
 143             rules = await self.conn.get_property('tokenizer_import_normalisation')
 144             return Transliterator.createFromRules("normalization", rules)
 145
 146         self.normalizer = await self.conn.get_cached_value('ICUTOK', 'normalizer',
 147                                                            _make_normalizer)
 148
 149         async def _make_transliterator() -> Any:
 150             rules = await self.conn.get_property('tokenizer_import_transliteration')
 151             return Transliterator.createFromRules("transliteration", rules)
 152
 153         self.transliterator = await self.conn.get_cached_value('ICUTOK', 'transliterator',
 154                                                                _make_transliterator)
 155
 156         if 'word' not in self.conn.t.meta.tables:
 157             sa.Table('word', self.conn.t.meta,
 158                      sa.Column('word_id', sa.Integer),
 159                      sa.Column('word_token', sa.Text, nullable=False),
 160                      sa.Column('type', sa.Text, nullable=False),
 161                      sa.Column('word', sa.Text),
 162                      sa.Column('info', self.conn.t.types.Json))
 163
 164
 165     async def analyze_query(self, phrases: List[qmod.Phrase]) -> qmod.QueryStruct:
 166         """ Analyze the given list of phrases and return the
 167             tokenized query.
 168         """
 169         log().section('Analyze query (using ICU tokenizer)')
 170         normalized = list(filter(lambda p: p.text,
 171                                  (qmod.Phrase(p.ptype, self.normalize_text(p.text))
 172                                   for p in phrases)))
 173         query = qmod.QueryStruct(normalized)
 174         log().var_dump('Normalized query', query.source)
 175         if not query.source:
 176             return query
 177
 178         parts, words = self.split_query(query)
 179         log().var_dump('Transliterated query', lambda: _dump_transliterated(query, parts))
 180
 181         for row in await self.lookup_in_db(list(words.keys())):
 182             for trange in words[row.word_token]:
 183                 token = ICUToken.from_db_row(row)
 184                 if row.type == 'S':
 185                     if row.info['op'] in ('in', 'near'):
 186                         if trange.start == 0:
 187                             query.add_token(trange, qmod.TokenType.NEAR_ITEM, token)
 188                     else:
 189                         query.add_token(trange, qmod.TokenType.QUALIFIER, token)
 190                         if trange.start == 0 or trange.end == query.num_token_slots():
 191                             token = copy(token)
 192                             token.penalty += 0.1 * (query.num_token_slots())
 193                             query.add_token(trange, qmod.TokenType.NEAR_ITEM, token)
 194                 else:
 195                     query.add_token(trange, DB_TO_TOKEN_TYPE[row.type], token)
 196
 197         self.add_extra_tokens(query, parts)
 198         self.rerank_tokens(query, parts)
 199
 200         log().table_dump('Word tokens', _dump_word_tokens(query))
 201
 202         return query
 203
 204
 205     def normalize_text(self, text: str) -> str:
 206         """ Bring the given text into a normalized form. That is the
 207             standardized form search will work with. All information removed
 208             at this stage is inevitably lost.
 209         """
 210         norm = cast(str, self.normalizer.transliterate(text))
 211         numspaces = norm.count(' ')
 212         if numspaces > 4 and len(norm) <= (numspaces + 1) * 3:
 213             return ''
 214
 215         return norm
 216
 217
 218     def split_query(self, query: qmod.QueryStruct) -> Tuple[QueryParts, WordDict]:
 219         """ Transliterate the phrases and split them into tokens.
 220
 221             Returns the list of transliterated tokens together with their
 222             normalized form and a dictionary of words for lookup together
 223             with their position.
 224         """
 225         parts: QueryParts = []
 226         phrase_start = 0
 227         words = defaultdict(list)
 228         wordnr = 0
 229         for phrase in query.source:
 230             query.nodes[-1].ptype = phrase.ptype
 231             for word in phrase.text.split(' '):
 232                 trans = self.transliterator.transliterate(word)
 233                 if trans:
 234                     for term in trans.split(' '):
 235                         if term:
 236                             parts.append(QueryPart(term, word, wordnr))
 237                             query.add_node(qmod.BreakType.TOKEN, phrase.ptype)
 238                     query.nodes[-1].btype = qmod.BreakType.WORD
 239                 wordnr += 1
 240             query.nodes[-1].btype = qmod.BreakType.PHRASE
 241
 242             for word, wrange in yield_words(parts, phrase_start):
 243                 words[word].append(wrange)
 244
 245             phrase_start = len(parts)
 246         query.nodes[-1].btype = qmod.BreakType.END
 247
 248         return parts, words
 249
 250
 251     async def lookup_in_db(self, words: List[str]) -> 'sa.Result[Any]':
 252         """ Return the token information from the database for the
 253             given word tokens.
 254         """
 255         t = self.conn.t.meta.tables['word']
 256         return await self.conn.execute(t.select().where(t.c.word_token.in_(words)))
 257
 258
 259     def add_extra_tokens(self, query: qmod.QueryStruct, parts: QueryParts) -> None:
 260         """ Add tokens to query that are not saved in the database.
 261         """
 262         for part, node, i in zip(parts, query.nodes, range(1000)):
 263             if len(part.token) <= 4 and part[0].isdigit()\
 264                and not node.has_tokens(i+1, qmod.TokenType.HOUSENUMBER):
 265                 query.add_token(qmod.TokenRange(i, i+1), qmod.TokenType.HOUSENUMBER,
 266                                 ICUToken(0.5, 0, 1, part.token, True, part.token, None))
 267
 268
 269     def rerank_tokens(self, query: qmod.QueryStruct, parts: QueryParts) -> None:
 270         """ Add penalties to tokens that depend on presence of other token.
 271         """
 272         for i, node, tlist in query.iter_token_lists():
 273             if tlist.ttype == qmod.TokenType.POSTCODE:
 274                 for repl in node.starting:
 275                     if repl.end == tlist.end and repl.ttype != qmod.TokenType.POSTCODE \
 276                        and (repl.ttype != qmod.TokenType.HOUSENUMBER
 277                             or len(tlist.tokens[0].lookup_word) > 4):
 278                         repl.add_penalty(0.39)
 279             elif tlist.ttype == qmod.TokenType.HOUSENUMBER \
 280                  and len(tlist.tokens[0].lookup_word) <= 3:
 281                 if any(c.isdigit() for c in tlist.tokens[0].lookup_word):
 282                     for repl in node.starting:
 283                         if repl.end == tlist.end and repl.ttype != qmod.TokenType.HOUSENUMBER:
 284                             repl.add_penalty(0.5 - tlist.tokens[0].penalty)
 285             elif tlist.ttype not in (qmod.TokenType.COUNTRY, qmod.TokenType.PARTIAL):
 286                 norm = parts[i].normalized
 287                 for j in range(i + 1, tlist.end):
 288                     if parts[j - 1].word_number != parts[j].word_number:
 289                         norm += '  ' + parts[j].normalized
 290                 for token in tlist.tokens:
 291                     cast(ICUToken, token).rematch(norm)
 292
 293
 294 def _dump_transliterated(query: qmod.QueryStruct, parts: QueryParts) -> str:
 295     out = query.nodes[0].btype.value
 296     for node, part in zip(query.nodes[1:], parts):
 297         out += part.token + node.btype.value
 298     return out
 299
 300
 301 def _dump_word_tokens(query: qmod.QueryStruct) -> Iterator[List[Any]]:
 302     yield ['type', 'token', 'word_token', 'lookup_word', 'penalty', 'count', 'info']
 303     for node in query.nodes:
 304         for tlist in node.starting:
 305             for token in tlist.tokens:
 306                 t = cast(ICUToken, token)
 307                 yield [tlist.ttype.name, t.token, t.word_token or '',
 308                        t.lookup_word or '', t.penalty, t.count, t.info]
 309
 310
 311 async def create_query_analyzer(conn: SearchConnection) -> AbstractQueryAnalyzer:
 312     """ Create and set up a new query analyzer for a database based
 313         on the ICU tokenizer.
 314     """
 315     out = ICUQueryAnalyzer(conn)
 316     await out.setup()
 317
 318     return out