iam-data-chatbot/query_normalizer.py at main · i2amparis/iam-data-chatbot · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
import re


PHRASE_SYNONYMS: dict[str, list[str]] = {
    "carbon dioxide": ["co2"],
    "greenhouse gas": ["ghg", "emissions"],
    "greenhouse gases": ["ghg", "emissions"],
    "solar pv": ["solar"],
    "photovoltaic": ["solar"],
    "photovoltaics": ["solar"],
    "wind power": ["wind"],
    "wind turbine": ["wind"],
    "wind turbines": ["wind"],
    "gross domestic product": ["gdp"],
    "time series": ["data", "query", "timeseries"],
    "data values": ["data", "query"],
}


TOKEN_SYNONYMS: dict[str, list[str]] = {
    "emission": ["emissions"],
    "emissions": ["emissions"],
    "electricity": ["electricity"],
    "power": ["electricity"],
    "pathway": ["scenario"],
    "pathways": ["scenario"],
    "country": ["region"],
    "countries": ["region"],
    "geography": ["region"],
    "geographies": ["region"],
    "location": ["region"],
    "locations": ["region"],
    "chart": ["plot"],
    "graph": ["plot"],
    "visualize": ["plot"],
    "visualise": ["plot"],
    "compare": ["comparison"],
    "versus": ["comparison"],
    "vs": ["comparison"],
    "data": ["data", "query"],
    "value": ["data", "query"],
    "values": ["data", "query"],
}


def normalize_query_text(text: str, expand_synonyms: bool = True) -> str:
    value = str(text or "").lower()
    value = value.replace("–", "-").replace("—", "-").replace("_", " ")
    value = re.sub(r"[-/]+", " ", value)
    value = re.sub(r"[^a-z0-9|.\s]+", " ", value)
    value = re.sub(r"\s+", " ", value).strip()

    if not expand_synonyms:
        return value

    additions: list[str] = []
    padded = f" {value} "
    for phrase, canonical_terms in PHRASE_SYNONYMS.items():
        if f" {phrase} " in padded:
            additions.extend(canonical_terms)

    for token in re.findall(r"[a-z0-9]+", value):
        additions.extend(TOKEN_SYNONYMS.get(token, []))
        if len(token) > 4 and token.endswith("s") and not token.endswith("ss"):
            additions.append(token[:-1])

    if additions:
        value = f"{value} {' '.join(additions)}"
    return re.sub(r"\s+", " ", value).strip()


def query_tokens(text: str) -> set[str]:
    normalized = normalize_query_text(text)
    return {token for token in re.findall(r"[a-z0-9]+", normalized) if token}