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, 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