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