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