Files
agent/app/modules/rag/intent_router_v2/normalization_terms.py
T

49 lines
1.9 KiB
Python

from __future__ import annotations
class KeyTermCanonicalizer:
_ALIASES: dict[str, set[str]] = {
"файл": {
"файл",
"файла",
"файле",
"файлу",
"файлом",
"файлы",
"файлов",
"файлам",
"файлами",
},
"класс": {"класс", "класса", "классе", "классу", "классом", "классы", "классов", "классам"},
"функция": {"функция", "функции", "функцию", "функцией", "функциях"},
"метод": {"метод", "метода", "методе", "методу", "методом", "методы"},
"документация": {"документация", "документации", "документацию"},
"тест": {"тест", "тесты", "тестов", "тестам", "тестами", "юнит-тест", "юниттест"},
"модуль": {"модуль", "модуля"},
"пакет": {"пакет"},
}
def __init__(self) -> None:
self._token_to_canonical = self._build_index()
def canonicalize(self, token: str) -> str | None:
return self._token_to_canonical.get((token or "").lower())
def aliases(self) -> set[str]:
values: set[str] = set()
for forms in self._ALIASES.values():
values.update(forms)
return values
def is_test_term(self, token: str) -> bool:
canonical = self.canonicalize(token)
return canonical == "тест"
def _build_index(self) -> dict[str, str]:
index: dict[str, str] = {}
for canonical, forms in self._ALIASES.items():
index[canonical] = canonical
for form in forms:
index[form] = canonical
return index