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, 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 ..typing import SaRow
  20 from ..sql.sqlalchemy_types import Json
  21 from ..connection import SearchConnection
  22 from ..logging import log
  23 from ..search import query as qmod
  24 from ..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         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,
 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         return cast(str, self.normalizer.transliterate(text))
 212
 213
 214     def split_query(self, query: qmod.QueryStruct) -> Tuple[QueryParts, WordDict]:
 215         """ Transliterate the phrases and split them into tokens.
 216
 217             Returns the list of transliterated tokens together with their
 218             normalized form and a dictionary of words for lookup together
 219             with their position.
 220         """
 221         parts: QueryParts = []
 222         phrase_start = 0
 223         words = defaultdict(list)
 224         wordnr = 0
 225         for phrase in query.source:
 226             query.nodes[-1].ptype = phrase.ptype
 227             for word in phrase.text.split(' '):
 228                 trans = self.transliterator.transliterate(word)
 229                 if trans:
 230                     for term in trans.split(' '):
 231                         if term:
 232                             parts.append(QueryPart(term, word, wordnr))
 233                             query.add_node(qmod.BreakType.TOKEN, phrase.ptype)
 234                     query.nodes[-1].btype = qmod.BreakType.WORD
 235                 wordnr += 1
 236             query.nodes[-1].btype = qmod.BreakType.PHRASE
 237
 238             for word, wrange in yield_words(parts, phrase_start):
 239                 words[word].append(wrange)
 240
 241             phrase_start = len(parts)
 242         query.nodes[-1].btype = qmod.BreakType.END
 243
 244         return parts, words
 245
 246
 247     async def lookup_in_db(self, words: List[str]) -> 'sa.Result[Any]':
 248         """ Return the token information from the database for the
 249             given word tokens.
 250         """
 251         t = self.conn.t.meta.tables['word']
 252         return await self.conn.execute(t.select().where(t.c.word_token.in_(words)))
 253
 254
 255     def add_extra_tokens(self, query: qmod.QueryStruct, parts: QueryParts) -> None:
 256         """ Add tokens to query that are not saved in the database.
 257         """
 258         for part, node, i in zip(parts, query.nodes, range(1000)):
 259             if len(part.token) <= 4 and part[0].isdigit()\
 260                and not node.has_tokens(i+1, qmod.TokenType.HOUSENUMBER):
 261                 query.add_token(qmod.TokenRange(i, i+1), qmod.TokenType.HOUSENUMBER,
 262                                 ICUToken(penalty=0.5, token=0,
 263                                          count=1, addr_count=1, lookup_word=part.token,
 264                                          word_token=part.token, info=None))
 265
 266
 267     def rerank_tokens(self, query: qmod.QueryStruct, parts: QueryParts) -> None:
 268         """ Add penalties to tokens that depend on presence of other token.
 269         """
 270         for i, node, tlist in query.iter_token_lists():
 271             if tlist.ttype == qmod.TokenType.POSTCODE:
 272                 for repl in node.starting:
 273                     if repl.end == tlist.end and repl.ttype != qmod.TokenType.POSTCODE \
 274                        and (repl.ttype != qmod.TokenType.HOUSENUMBER
 275                             or len(tlist.tokens[0].lookup_word) > 4):
 276                         repl.add_penalty(0.39)
 277             elif tlist.ttype == qmod.TokenType.HOUSENUMBER \
 278                  and len(tlist.tokens[0].lookup_word) <= 3:
 279                 if any(c.isdigit() for c in tlist.tokens[0].lookup_word):
 280                     for repl in node.starting:
 281                         if repl.end == tlist.end and repl.ttype != qmod.TokenType.HOUSENUMBER:
 282                             repl.add_penalty(0.5 - tlist.tokens[0].penalty)
 283             elif tlist.ttype not in (qmod.TokenType.COUNTRY, qmod.TokenType.PARTIAL):
 284                 norm = parts[i].normalized
 285                 for j in range(i + 1, tlist.end):
 286                     if parts[j - 1].word_number != parts[j].word_number:
 287                         norm += '  ' + parts[j].normalized
 288                 for token in tlist.tokens:
 289                     cast(ICUToken, token).rematch(norm)
 290
 291
 292 def _dump_transliterated(query: qmod.QueryStruct, parts: QueryParts) -> str:
 293     out = query.nodes[0].btype.value
 294     for node, part in zip(query.nodes[1:], parts):
 295         out += part.token + node.btype.value
 296     return out
 297
 298
 299 def _dump_word_tokens(query: qmod.QueryStruct) -> Iterator[List[Any]]:
 300     yield ['type', 'token', 'word_token', 'lookup_word', 'penalty', 'count', 'info']
 301     for node in query.nodes:
 302         for tlist in node.starting:
 303             for token in tlist.tokens:
 304                 t = cast(ICUToken, token)
 305                 yield [tlist.ttype.name, 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     out = ICUQueryAnalyzer(conn)
 314     await out.setup()
 315
 316     return out