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 from nominatim.db.sqlalchemy_types import Json
  26
  27
  28 DB_TO_TOKEN_TYPE = {
  29     'W': qmod.TokenType.WORD,
  30     'w': qmod.TokenType.PARTIAL,
  31     'H': qmod.TokenType.HOUSENUMBER,
  32     'P': qmod.TokenType.POSTCODE,
  33     'C': qmod.TokenType.COUNTRY
  34 }
  35
  36
  37 class QueryPart(NamedTuple):
  38     """ Normalized and transliterated form of a single term in the query.
  39         When the term came out of a split during the transliteration,
  40         the normalized string is the full word before transliteration.
  41         The word number keeps track of the word before transliteration
  42         and can be used to identify partial transliterated terms.
  43     """
  44     token: str
  45     normalized: str
  46     word_number: int
  47
  48
  49 QueryParts = List[QueryPart]
  50 WordDict = Dict[str, List[qmod.TokenRange]]
  51
  52 def yield_words(terms: List[QueryPart], start: int) -> Iterator[Tuple[str, qmod.TokenRange]]:
  53     """ Return all combinations of words in the terms list after the
  54         given position.
  55     """
  56     total = len(terms)
  57     for first in range(start, total):
  58         word = terms[first].token
  59         yield word, qmod.TokenRange(first, first + 1)
  60         for last in range(first + 1, min(first + 20, total)):
  61             word = ' '.join((word, terms[last].token))
  62             yield word, qmod.TokenRange(first, last + 1)
  63
  64
  65 @dataclasses.dataclass
  66 class ICUToken(qmod.Token):
  67     """ Specialised token for ICU tokenizer.
  68     """
  69     word_token: str
  70     info: Optional[Dict[str, Any]]
  71
  72     def get_category(self) -> Tuple[str, str]:
  73         assert self.info
  74         return self.info.get('class', ''), self.info.get('type', '')
  75
  76
  77     def rematch(self, norm: str) -> None:
  78         """ Check how well the token matches the given normalized string
  79             and add a penalty, if necessary.
  80         """
  81         if not self.lookup_word:
  82             return
  83
  84         seq = difflib.SequenceMatcher(a=self.lookup_word, b=norm)
  85         distance = 0
  86         for tag, afrom, ato, bfrom, bto in seq.get_opcodes():
  87             if tag in ('delete', 'insert') and (afrom == 0 or ato == len(self.lookup_word)):
  88                 distance += 1
  89             elif tag == 'replace':
  90                 distance += max((ato-afrom), (bto-bfrom))
  91             elif tag != 'equal':
  92                 distance += abs((ato-afrom) - (bto-bfrom))
  93         self.penalty += (distance/len(self.lookup_word))
  94
  95
  96     @staticmethod
  97     def from_db_row(row: SaRow) -> 'ICUToken':
  98         """ Create a ICUToken from the row of the word table.
  99         """
 100         count = 1 if row.info is None else row.info.get('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=count,
 126                         lookup_word=lookup_word, is_indexed=True,
 127                         word_token=row.word_token, info=row.info)
 128
 129
 130
 131 class ICUQueryAnalyzer(AbstractQueryAnalyzer):
 132     """ Converter for query strings into a tokenized query
 133         using the tokens created by a ICU tokenizer.
 134     """
 135
 136     def __init__(self, conn: SearchConnection) -> None:
 137         self.conn = conn
 138
 139
 140     async def setup(self) -> None:
 141         """ Set up static data structures needed for the analysis.
 142         """
 143         async def _make_normalizer() -> Any:
 144             rules = await self.conn.get_property('tokenizer_import_normalisation')
 145             return Transliterator.createFromRules("normalization", rules)
 146
 147         self.normalizer = await self.conn.get_cached_value('ICUTOK', 'normalizer',
 148                                                            _make_normalizer)
 149
 150         async def _make_transliterator() -> Any:
 151             rules = await self.conn.get_property('tokenizer_import_transliteration')
 152             return Transliterator.createFromRules("transliteration", rules)
 153
 154         self.transliterator = await self.conn.get_cached_value('ICUTOK', 'transliterator',
 155                                                                _make_transliterator)
 156
 157         if 'word' not in self.conn.t.meta.tables:
 158             sa.Table('word', self.conn.t.meta,
 159                      sa.Column('word_id', sa.Integer),
 160                      sa.Column('word_token', sa.Text, nullable=False),
 161                      sa.Column('type', sa.Text, nullable=False),
 162                      sa.Column('word', sa.Text),
 163                      sa.Column('info', Json))
 164
 165
 166     async def analyze_query(self, phrases: List[qmod.Phrase]) -> qmod.QueryStruct:
 167         """ Analyze the given list of phrases and return the
 168             tokenized query.
 169         """
 170         log().section('Analyze query (using ICU tokenizer)')
 171         normalized = list(filter(lambda p: p.text,
 172                                  (qmod.Phrase(p.ptype, self.normalize_text(p.text))
 173                                   for p in phrases)))
 174         query = qmod.QueryStruct(normalized)
 175         log().var_dump('Normalized query', query.source)
 176         if not query.source:
 177             return query
 178
 179         parts, words = self.split_query(query)
 180         log().var_dump('Transliterated query', lambda: _dump_transliterated(query, parts))
 181
 182         for row in await self.lookup_in_db(list(words.keys())):
 183             for trange in words[row.word_token]:
 184                 token = ICUToken.from_db_row(row)
 185                 if row.type == 'S':
 186                     if row.info['op'] in ('in', 'near'):
 187                         if trange.start == 0:
 188                             query.add_token(trange, qmod.TokenType.NEAR_ITEM, token)
 189                     else:
 190                         query.add_token(trange, qmod.TokenType.QUALIFIER, token)
 191                         if trange.start == 0 or trange.end == query.num_token_slots():
 192                             token = copy(token)
 193                             token.penalty += 0.1 * (query.num_token_slots())
 194                             query.add_token(trange, qmod.TokenType.NEAR_ITEM, 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, 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