src/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) 2024 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, Iterator, Any, cast
  11 import dataclasses
  12 import difflib
  13 import re
  14 from itertools import zip_longest
  15
  16 from icu import Transliterator
  17
  18 import sqlalchemy as sa
  19
  20 from ..errors import UsageError
  21 from ..typing import SaRow
  22 from ..sql.sqlalchemy_types import Json
  23 from ..connection import SearchConnection
  24 from ..logging import log
  25 from . import query as qmod
  26 from ..query_preprocessing.config import QueryConfig
  27 from ..query_preprocessing.base import QueryProcessingFunc
  28 from .query_analyzer_factory import AbstractQueryAnalyzer
  29 from .postcode_parser import PostcodeParser
  30
  31
  32 DB_TO_TOKEN_TYPE = {
  33     'W': qmod.TOKEN_WORD,
  34     'w': qmod.TOKEN_PARTIAL,
  35     'H': qmod.TOKEN_HOUSENUMBER,
  36     'P': qmod.TOKEN_POSTCODE,
  37     'C': qmod.TOKEN_COUNTRY
  38 }
  39
  40 PENALTY_IN_TOKEN_BREAK = {
  41      qmod.BREAK_START: 0.5,
  42      qmod.BREAK_END: 0.5,
  43      qmod.BREAK_PHRASE: 0.5,
  44      qmod.BREAK_SOFT_PHRASE: 0.5,
  45      qmod.BREAK_WORD: 0.1,
  46      qmod.BREAK_PART: 0.0,
  47      qmod.BREAK_TOKEN: 0.0
  48 }
  49
  50
  51 @dataclasses.dataclass
  52 class ICUToken(qmod.Token):
  53     """ Specialised token for ICU tokenizer.
  54     """
  55     word_token: str
  56     info: Optional[Dict[str, Any]]
  57
  58     def get_category(self) -> Tuple[str, str]:
  59         assert self.info
  60         return self.info.get('class', ''), self.info.get('type', '')
  61
  62     def rematch(self, norm: str) -> None:
  63         """ Check how well the token matches the given normalized string
  64             and add a penalty, if necessary.
  65         """
  66         if not self.lookup_word:
  67             return
  68
  69         seq = difflib.SequenceMatcher(a=self.lookup_word, b=norm)
  70         distance = 0
  71         for tag, afrom, ato, bfrom, bto in seq.get_opcodes():
  72             if tag in ('delete', 'insert') and (afrom == 0 or ato == len(self.lookup_word)):
  73                 distance += 1
  74             elif tag == 'replace':
  75                 distance += max((ato-afrom), (bto-bfrom))
  76             elif tag != 'equal':
  77                 distance += abs((ato-afrom) - (bto-bfrom))
  78         self.penalty += (distance/len(self.lookup_word))
  79
  80     @staticmethod
  81     def from_db_row(row: SaRow, base_penalty: float = 0.0) -> 'ICUToken':
  82         """ Create a ICUToken from the row of the word table.
  83         """
  84         count = 1 if row.info is None else row.info.get('count', 1)
  85         addr_count = 1 if row.info is None else row.info.get('addr_count', 1)
  86
  87         penalty = base_penalty
  88         if row.type == 'w':
  89             penalty += 0.3
  90         elif row.type == 'W':
  91             if len(row.word_token) == 1 and row.word_token == row.word:
  92                 penalty += 0.2 if row.word.isdigit() else 0.3
  93         elif row.type == 'H':
  94             penalty += sum(0.1 for c in row.word_token if c != ' ' and not c.isdigit())
  95             if all(not c.isdigit() for c in row.word_token):
  96                 penalty += 0.2 * (len(row.word_token) - 1)
  97         elif row.type == 'C':
  98             if len(row.word_token) == 1:
  99                 penalty += 0.3
 100
 101         if row.info is None:
 102             lookup_word = row.word
 103         else:
 104             lookup_word = row.info.get('lookup', row.word)
 105         if lookup_word:
 106             lookup_word = lookup_word.split('@', 1)[0]
 107         else:
 108             lookup_word = row.word_token
 109
 110         return ICUToken(penalty=penalty, token=row.word_id, count=max(1, count),
 111                         lookup_word=lookup_word,
 112                         word_token=row.word_token, info=row.info,
 113                         addr_count=max(1, addr_count))
 114
 115
 116 @dataclasses.dataclass
 117 class ICUAnalyzerConfig:
 118     postcode_parser: PostcodeParser
 119     normalizer: Transliterator
 120     transliterator: Transliterator
 121     preprocessors: List[QueryProcessingFunc]
 122
 123     @staticmethod
 124     async def create(conn: SearchConnection) -> 'ICUAnalyzerConfig':
 125         rules = await conn.get_property('tokenizer_import_normalisation')
 126         normalizer = Transliterator.createFromRules("normalization", rules)
 127
 128         rules = await conn.get_property('tokenizer_import_transliteration')
 129         transliterator = Transliterator.createFromRules("transliteration", rules)
 130
 131         preprocessing_rules = conn.config.load_sub_configuration('icu_tokenizer.yaml',
 132                                                                  config='TOKENIZER_CONFIG')\
 133                                          .get('query-preprocessing', [])
 134
 135         preprocessors: List[QueryProcessingFunc] = []
 136         for func in preprocessing_rules:
 137             if 'step' not in func:
 138                 raise UsageError("Preprocessing rule is missing the 'step' attribute.")
 139             if not isinstance(func['step'], str):
 140                 raise UsageError("'step' attribute must be a simple string.")
 141
 142             module = conn.config.load_plugin_module(
 143                         func['step'], 'nominatim_api.query_preprocessing')
 144             preprocessors.append(
 145                 module.create(QueryConfig(func).set_normalizer(normalizer)))
 146
 147         return ICUAnalyzerConfig(PostcodeParser(conn.config),
 148                                  normalizer, transliterator, preprocessors)
 149
 150
 151 class ICUQueryAnalyzer(AbstractQueryAnalyzer):
 152     """ Converter for query strings into a tokenized query
 153         using the tokens created by a ICU tokenizer.
 154     """
 155     def __init__(self, conn: SearchConnection, config: ICUAnalyzerConfig) -> None:
 156         self.conn = conn
 157         self.postcode_parser = config.postcode_parser
 158         self.normalizer = config.normalizer
 159         self.transliterator = config.transliterator
 160         self.preprocessors = config.preprocessors
 161
 162     async def analyze_query(self, phrases: List[qmod.Phrase]) -> qmod.QueryStruct:
 163         """ Analyze the given list of phrases and return the
 164             tokenized query.
 165         """
 166         log().section('Analyze query (using ICU tokenizer)')
 167         for func in self.preprocessors:
 168             phrases = func(phrases)
 169
 170         if len(phrases) == 1 \
 171                 and phrases[0].text.count(' ') > 3 \
 172                 and max(len(s) for s in phrases[0].text.split()) < 3:
 173             normalized = []
 174
 175         query = qmod.QueryStruct(phrases)
 176
 177         log().var_dump('Normalized query', query.source)
 178         if not query.source:
 179             return query
 180
 181         self.split_query(query)
 182         log().var_dump('Transliterated query', lambda: query.get_transliterated_query())
 183         words = query.extract_words(base_penalty=PENALTY_IN_TOKEN_BREAK[qmod.BREAK_WORD])
 184
 185         for row in await self.lookup_in_db(list(words.keys())):
 186             for trange in words[row.word_token]:
 187                 token = ICUToken.from_db_row(row, trange.penalty or 0.0)
 188                 if row.type == 'S':
 189                     if row.info['op'] in ('in', 'near'):
 190                         if trange.start == 0:
 191                             query.add_token(trange, qmod.TOKEN_NEAR_ITEM, token)
 192                     else:
 193                         if trange.start == 0 and trange.end == query.num_token_slots():
 194                             query.add_token(trange, qmod.TOKEN_NEAR_ITEM, token)
 195                         else:
 196                             query.add_token(trange, qmod.TOKEN_QUALIFIER, token)
 197                 else:
 198                     query.add_token(trange, DB_TO_TOKEN_TYPE[row.type], token)
 199
 200         self.add_extra_tokens(query)
 201         for start, end, pc in self.postcode_parser.parse(query):
 202             term = ' '.join(n.term_lookup for n in query.nodes[start + 1:end + 1])
 203             query.add_token(qmod.TokenRange(start, end),
 204                             qmod.TOKEN_POSTCODE,
 205                             ICUToken(penalty=0.1, token=0, count=1, addr_count=1,
 206                                      lookup_word=pc, word_token=term,
 207                                      info=None))
 208         self.rerank_tokens(query)
 209
 210         log().table_dump('Word tokens', _dump_word_tokens(query))
 211
 212         return query
 213
 214     def normalize_text(self, text: str) -> str:
 215         """ Bring the given text into a normalized form. That is the
 216             standardized form search will work with. All information removed
 217             at this stage is inevitably lost.
 218         """
 219         return cast(str, self.normalizer.transliterate(text)).strip('-: ')
 220
 221     def split_query(self, query: qmod.QueryStruct) -> None:
 222         """ Transliterate the phrases and split them into tokens.
 223         """
 224         for phrase in query.source:
 225             query.nodes[-1].ptype = phrase.ptype
 226             phrase_split = re.split('([ :-])', phrase.text)
 227             # The zip construct will give us the pairs of word/break from
 228             # the regular expression split. As the split array ends on the
 229             # final word, we simply use the fillvalue to even out the list and
 230             # add the phrase break at the end.
 231             for word, breakchar in zip_longest(*[iter(phrase_split)]*2, fillvalue=','):
 232                 if not word:
 233                     continue
 234                 trans = self.transliterator.transliterate(word)
 235                 if trans:
 236                     for term in trans.split(' '):
 237                         if term:
 238                             query.add_node(qmod.BREAK_TOKEN, phrase.ptype,
 239                                            PENALTY_IN_TOKEN_BREAK[qmod.BREAK_TOKEN],
 240                                            term, word)
 241                     query.nodes[-1].adjust_break(breakchar,
 242                                                  PENALTY_IN_TOKEN_BREAK[breakchar])
 243
 244         query.nodes[-1].adjust_break(qmod.BREAK_END, PENALTY_IN_TOKEN_BREAK[qmod.BREAK_END])
 245
 246     async def lookup_in_db(self, words: List[str]) -> 'sa.Result[Any]':
 247         """ Return the token information from the database for the
 248             given word tokens.
 249
 250             This function excludes postcode tokens
 251         """
 252         t = self.conn.t.meta.tables['word']
 253         return await self.conn.execute(t.select()
 254                                         .where(t.c.word_token.in_(words))
 255                                         .where(t.c.type != 'P'))
 256
 257     def add_extra_tokens(self, query: qmod.QueryStruct) -> None:
 258         """ Add tokens to query that are not saved in the database.
 259         """
 260         need_hnr = False
 261         for i, node in enumerate(query.nodes):
 262             is_full_token = node.btype not in (qmod.BREAK_TOKEN, qmod.BREAK_PART)
 263             if need_hnr and is_full_token \
 264                     and len(node.term_normalized) <= 4 and node.term_normalized.isdigit():
 265                 query.add_token(qmod.TokenRange(i-1, i), qmod.TOKEN_HOUSENUMBER,
 266                                 ICUToken(penalty=0.5, token=0,
 267                                          count=1, addr_count=1,
 268                                          lookup_word=node.term_lookup,
 269                                          word_token=node.term_lookup, info=None))
 270
 271             need_hnr = is_full_token and not node.has_tokens(i+1, qmod.TOKEN_HOUSENUMBER)
 272
 273     def rerank_tokens(self, query: qmod.QueryStruct) -> None:
 274         """ Add penalties to tokens that depend on presence of other token.
 275         """
 276         for i, node, tlist in query.iter_token_lists():
 277             if tlist.ttype == qmod.TOKEN_POSTCODE:
 278                 tlen = len(cast(ICUToken, tlist.tokens[0]).word_token)
 279                 for repl in node.starting:
 280                     if repl.end == tlist.end and repl.ttype != qmod.TOKEN_POSTCODE \
 281                        and (repl.ttype != qmod.TOKEN_HOUSENUMBER or tlen > 4):
 282                         repl.add_penalty(0.39)
 283             elif (tlist.ttype == qmod.TOKEN_HOUSENUMBER
 284                   and len(tlist.tokens[0].lookup_word) <= 3):
 285                 if any(c.isdigit() for c in tlist.tokens[0].lookup_word):
 286                     for repl in node.starting:
 287                         if repl.end == tlist.end and repl.ttype != qmod.TOKEN_HOUSENUMBER:
 288                             repl.add_penalty(0.5 - tlist.tokens[0].penalty)
 289             elif tlist.ttype not in (qmod.TOKEN_COUNTRY, qmod.TOKEN_PARTIAL):
 290                 norm = ' '.join(n.term_normalized for n in query.nodes[i + 1:tlist.end + 1]
 291                                 if n.btype != qmod.BREAK_TOKEN)
 292                 if not norm:
 293                     # Can happen when the token only covers a partial term
 294                     norm = query.nodes[i + 1].term_normalized
 295                 for token in tlist.tokens:
 296                     cast(ICUToken, token).rematch(norm)
 297
 298
 299 def _dump_word_tokens(query: qmod.QueryStruct) -> Iterator[List[Any]]:
 300     yield ['type', 'from', 'to', 'token', 'word_token', 'lookup_word', 'penalty', 'count', 'info']
 301     for i, node in enumerate(query.nodes):
 302         for tlist in node.starting:
 303             for token in tlist.tokens:
 304                 t = cast(ICUToken, token)
 305                 yield [tlist.ttype, str(i), str(tlist.end), t.token, t.word_token or '',
 306                        t.lookup_word or '', t.penalty, t.count, t.info]
 307
 308
 309 async def create_query_analyzer(conn: SearchConnection) -> AbstractQueryAnalyzer:
 310     """ Create and set up a new query analyzer for a database based
 311         on the ICU tokenizer.
 312     """
 313     async def _get_config() -> ICUAnalyzerConfig:
 314         if 'word' not in conn.t.meta.tables:
 315             sa.Table('word', conn.t.meta,
 316                      sa.Column('word_id', sa.Integer),
 317                      sa.Column('word_token', sa.Text, nullable=False),
 318                      sa.Column('type', sa.Text, nullable=False),
 319                      sa.Column('word', sa.Text),
 320                      sa.Column('info', Json))
 321
 322         return await ICUAnalyzerConfig.create(conn)
 323
 324     config = await conn.get_cached_value('ICUTOK', 'config', _get_config)
 325
 326     return ICUQueryAnalyzer(conn, config)