diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000000000000000000000000000000000000..701a8856a2ee85cba6a9f575d2f94e5ff03797bc --- /dev/null +++ b/.dockerignore @@ -0,0 +1,5 @@ +/data/ +/logs/ +__pycache__ +*.db +.env \ No newline at end of file diff --git a/.env-example b/.env-example new file mode 100644 index 0000000000000000000000000000000000000000..fc8eef76d1f84590e31d9851836a78d970785e01 --- /dev/null +++ b/.env-example @@ -0,0 +1 @@ +DEEPINFRA_API_KEY=Bearer <ключ> \ No newline at end of file diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..681a32746053c4dd5544fc3301df570927dc8e1d --- /dev/null +++ b/.gitignore @@ -0,0 +1,19 @@ +/data/ +common.log +/output/test.json +/logs/ + +venv +.idea +__pycache__ +*.db + +*.docx +*.doc +*.pdf +*.xlsx +*.xls +*.pptx +*.ppt +.env +/docker-compose.yaml \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..d1ab9758fa78ad96f9bb3bc6057b06305363df50 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,44 @@ +FROM nvidia/cuda:12.6.0-runtime-ubuntu22.04 + +ARG PORT=7860 +ENV PORT=${PORT} +ENV CONFIG_PATH=config_dev.yaml +ENV SQLALCHEMY_DATABASE_URL=sqlite:///./logs.db + +ENV PYTHONUNBUFFERED=1 +ENV DEBIAN_FRONTEND=noninteractive + +WORKDIR /app + +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + python3.11 \ + python3.11-distutils \ + wget \ + && wget https://bootstrap.pypa.io/get-pip.py \ + && python3.11 get-pip.py \ + && rm get-pip.py \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* + +# Set Python 3.11 as the default python3 +RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1 \ + && update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1 + +# Устанавливаем специфичные версии библиотек PyTorch +RUN python -m pip install \ + torch==2.6.0+cu126 \ + --index-url https://download.pytorch.org/whl/cu126 + +COPY requirements.txt /app/ +RUN python -m pip install -r requirements.txt +# RUN python -m pip install --ignore-installed elasticsearch==7.11.0 || true + +COPY . . + +RUN mkdir -p /app/data/regulation_datasets /app/data/documents /app/logs + + +EXPOSE ${PORT} + +CMD ["sh", "-c", "uvicorn main:app --host 0.0.0.0 --port ${PORT}"] diff --git a/README.md b/README.md index 690d54ed68f153d12b6a983e59e571403457c73c..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 100644 --- a/README.md +++ b/README.md @@ -1,10 +0,0 @@ ---- -title: Generic Chatbot Backend -emoji: 🚀 -colorFrom: pink -colorTo: gray -sdk: docker -pinned: false ---- - -Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference diff --git a/common/common.py b/common/common.py new file mode 100644 index 0000000000000000000000000000000000000000..6f3e75fb54bd5570c3b400f7a297b77fe8e84c88 --- /dev/null +++ b/common/common.py @@ -0,0 +1,235 @@ +import logging +from enum import Enum + + +def configure_logging(level=logging.INFO, config_file_path='./common.log'): + logging.basicConfig( + filename=config_file_path, + filemode="a", + level=level, + datefmt="%Y-%m-%d %H:%M:%S", + format="[%(asctime)s.%(msecs)03d] %(module)30s:%(lineno)4d %(levelname)-7s - %(message)s", + ) + + +def get_elastic_query(query): + return { + "query": { + "multi_match": { + "query": f"{query}", + "fields": ["text"], + "fuzziness": "AUTO", + "analyzer": "russian", + } + } + } + + +def get_elastic_people_query(query): + has_business_curator = ( + "бизнес куратор" in query.lower() + or "бизнес-куратор" in query.lower() + or "куратор" in query.lower() + ) + business_curator_boost = 30 if has_business_curator else 15 + return { + "query": { + "bool": { + "should": [ + { + "multi_match": { + "query": f"{query}", + "fields": ["person_name^3"], + "fuzziness": "AUTO", + "analyzer": "standard", + } + }, + { + "nested": { + "path": "business_processes", + "query": { + "multi_match": { + "query": f"{query}", + "fields": [ + "business_processes.production_activities_section", + "business_processes.processes_name", + ], + "fuzziness": "AUTO", + "analyzer": "standard", + } + }, + } + }, + { + "nested": { + "path": "organizatinal_structure", + "query": { + "multi_match": { + "query": f"{query}", + "fields": ["organizatinal_structure.position^2"], + "fuzziness": "AUTO", + "analyzer": "standard", + } + }, + } + }, + { + "nested": { + "path": "business_curator", + "query": { + "multi_match": { + "query": f"{query}", + "fields": [ + f"business_curator.company_name^{business_curator_boost}" + ], + "fuzziness": "AUTO", + "analyzer": "standard", + } + }, + } + }, + ] + } + }, + "min_score": 13.0, + } + + +def get_elastic_group_query(query): + return { + "query": { + "bool": { + "should": [ + { + "multi_match": { + "query": f"{query}", + "fields": ["group_name"], + "fuzziness": "AUTO", + "analyzer": "standard", + } + }, + { + "multi_match": { + "query": "персонального состава Персональный состав Комитета ПАО ГМК Норильский никель Рабочей группы", + "fields": ["group_name"], + "operator": "or", + "boost": 0.1, + } + }, + ] + } + }, + "min_score": 7.5, + } + + +def get_elastic_rocks_nn_query(query): + return { + "query": { + "function_score": { + "query": { + "multi_match": { + "query": f"{query}", + "fields": ["division_name", "division_name_2", "company_name"], + "fuzziness": "AUTO", + "analyzer": "custom_analyzer", + } + }, + "functions": [{"filter": {"term": {"_id": "3"}}, "weight": 0.5}], + "boost_mode": "multiply", + } + }, + "min_score": 0.5, + } + + +def get_elastic_segmentation_query(query): + return { + "query": { + "bool": { + "should": [ + { + "multi_match": { + "query": f"{query}", + "fields": [ + "segmentation_model", + "segmentation_model2", + "company_name", + ], + "fuzziness": "AUTO", + "analyzer": "russian", + } + }, + { + "multi_match": { + "query": "модели сегментации модель сегментации", + "fields": ["segmentation_model", "segmentation_model2"], + "operator": "or", + "boost": 0.1, + } + }, + ] + } + }, + "min_score": 1.0, + } + + +def get_elastic_abbreviation_query(query): + return { + "query": { + "multi_match": { + "query": f"{query}", + "fuzziness": "AUTO", + "fields": ["text"], + "analyzer": "russian", + } + } + } + + +def combine_answer(answer): + """ + + Args: + answer: + + Returns: + + """ + answer_combined = {} + indexes = [] + for key in answer: + if key != 'people_search': + for answer_key in answer[key]: + answer_value = answer[key][answer_key] + filename_i = answer_value["doc_name"] + title_i = answer_value["title"] + + if ( + filename_i in answer_combined + and answer_value['index_answer'] not in indexes + ): + answer_combined[filename_i]["chunks"].append(answer_value) + else: + answer_combined[filename_i] = { + "filename": filename_i, + "title": title_i, + "chunks": [answer_value], + } + indexes.append(answer_value['index_answer']) + return list(answer_combined.values()) + + +class TypeQuestion(Enum): + TYPE_ONE = '[1]' + TYPE_TWO = '[2]' + TYPE_THREE = '[3]' + + +def get_source_format(filename: str) -> str: + """ + Получает формат файла из имени файла. + """ + format_ = filename.split('.')[-1] + return format_.upper() diff --git a/common/configuration.py b/common/configuration.py new file mode 100644 index 0000000000000000000000000000000000000000..f31d0c1e4ab8fb6d8c66b8e6ef15412ea2f9feb0 --- /dev/null +++ b/common/configuration.py @@ -0,0 +1,274 @@ +"""This module includes classes to define configurations.""" + +from typing import Any, Dict, List, Optional + +from pyaml_env import parse_config +from pydantic import BaseModel + + +class Query(BaseModel): + query: str + query_abbreviation: str + abbreviations_replaced: Optional[List] = None + userName: Optional[str] = None + + +class SemanticChunk(BaseModel): + index_answer: int + doc_name: str + title: str + text_answer: str + # doc_number: str # TODO Потом поменять название переменной на doc_id везде с чем это будет связанно + other_info: List + start_index_paragraph: int + + +class FilterChunks(BaseModel): + id: str + filename: str + title: str + chunks: List[SemanticChunk] + + +class BusinessProcess(BaseModel): + production_activities_section: Optional[str] + processes_name: Optional[str] + level_process: Optional[str] + + +class Lead(BaseModel): + person: Optional[str] + leads: Optional[str] + + +class Subordinate(BaseModel): + person_name: Optional[str] + position: Optional[str] + + +class OrganizationalStructure(BaseModel): + position: Optional[str] = None + leads: Optional[List[Lead]] = None + subordinates: Optional[Subordinate] = None + + +class RocksNN(BaseModel): + division: Optional[str] + company_name: Optional[str] + + +class RocksNNSearch(BaseModel): + division: Optional[str] + company_name: Optional[List] + + +class SegmentationSearch(BaseModel): + segmentation_model: Optional[str] + company_name: Optional[List] + + +class Group(BaseModel): + group_name: Optional[str] + position_in_group: Optional[str] + block: Optional[str] + + +class GroupComposition(BaseModel): + person_name: Optional[str] + position_in_group: Optional[str] + + +class SearchGroupComposition(BaseModel): + group_name: Optional[str] + group_composition: Optional[List[GroupComposition]] + + +class PeopleChunks(BaseModel): + business_processes: Optional[List[BusinessProcess]] = None + organizatinal_structure: Optional[List[OrganizationalStructure]] = None + business_curator: Optional[List[RocksNN]] = None + groups: Optional[List[Group]] = None + person_name: str + + +class SummaryChunks(BaseModel): + doc_chunks: Optional[List[FilterChunks]] = None + people_search: Optional[List[PeopleChunks]] = None + groups_search: Optional[SearchGroupComposition] = None + rocks_nn_search: Optional[RocksNNSearch] = None + segmentation_search: Optional[SegmentationSearch] = None + query_type: str = '[3]' + + +class ElasticConfiguration: + def __init__(self, config_data): + self.es_host = str(config_data['es_host']) + self.es_port = int(config_data['es_port']) + self.use_elastic = bool(config_data['use_elastic']) + self.people_path = str(config_data['people_path']) + + +class FaissDataConfiguration: + def __init__(self, config_data): + self.model_embedding_path = str(config_data['model_embedding_path']) + self.device = str(config_data['device']) + self.path_to_metadata = str(config_data['path_to_metadata']) + + +class ChunksElasticSearchConfiguration: + def __init__(self, config_data): + self.use_chunks_search = bool(config_data['use_chunks_search']) + self.index_name = str(config_data['index_name']) + self.k_neighbors = int(config_data['k_neighbors']) + + +class PeopleSearchConfiguration: + def __init__(self, config_data): + self.use_people_search = bool(config_data['use_people_search']) + self.index_name = str(config_data['index_name']) + self.k_neighbors = int(config_data['k_neighbors']) + + +class VectorSearchConfiguration: + def __init__(self, config_data): + self.use_vector_search = bool(config_data['use_vector_search']) + self.k_neighbors = int(config_data['k_neighbors']) + + +class GroupsSearchConfiguration: + def __init__(self, config_data): + self.use_groups_search = bool(config_data['use_groups_search']) + self.index_name = str(config_data['index_name']) + self.k_neighbors = int(config_data['k_neighbors']) + + +class RocksNNSearchConfiguration: + def __init__(self, config_data): + self.use_rocks_nn_search = bool(config_data['use_rocks_nn_search']) + self.index_name = str(config_data['index_name']) + self.k_neighbors = int(config_data['k_neighbors']) + + +class AbbreviationSearchConfiguration: + def __init__(self, config_data): + self.use_abbreviation_search = bool(config_data['use_abbreviation_search']) + self.index_name = str(config_data['index_name']) + self.k_neighbors = int(config_data['k_neighbors']) + + +class SegmentationSearchConfiguration: + def __init__(self, config_data): + self.use_segmentation_search = bool(config_data['use_segmentation_search']) + self.index_name = str(config_data['index_name']) + self.k_neighbors = int(config_data['k_neighbors']) + + +class SearchConfiguration: + def __init__(self, config_data): + self.vector_search = VectorSearchConfiguration(config_data['vector_search']) + self.people_elastic_search = PeopleSearchConfiguration( + config_data['people_elastic_search'] + ) + self.chunks_elastic_search = ChunksElasticSearchConfiguration( + config_data['chunks_elastic_search'] + ) + self.groups_elastic_search = GroupsSearchConfiguration( + config_data['groups_elastic_search'] + ) + self.rocks_nn_elastic_search = RocksNNSearchConfiguration( + config_data['rocks_nn_elastic_search'] + ) + self.segmentation_elastic_search = SegmentationSearchConfiguration( + config_data['segmentation_elastic_search'] + ) + self.stop_index_names = list(config_data['stop_index_names']) + self.abbreviation_search = AbbreviationSearchConfiguration( + config_data['abbreviation_search'] + ) + + +class FilesConfiguration: + def __init__(self, config_data): + self.empty_start = bool(config_data['empty_start']) + self.regulations_path = str(config_data['regulations_path']) + self.default_regulations_path = str(config_data['default_regulations_path']) + self.documents_path = str(config_data['documents_path']) + + +class RankingConfiguration: + def __init__(self, config_data): + self.use_ranging = bool(config_data['use_ranging']) + self.alpha = float(config_data['alpha']) + self.beta = float(config_data['beta']) + self.k_neighbors = int(config_data['k_neighbors']) + + +class DataBaseConfiguration: + def __init__(self, config_data): + self.elastic = ElasticConfiguration(config_data['elastic']) + self.faiss = FaissDataConfiguration(config_data['faiss']) + self.search = SearchConfiguration(config_data['search']) + self.files = FilesConfiguration(config_data['files']) + self.ranker = RankingConfiguration(config_data['ranging']) + + +class LLMConfiguration: + def __init__(self, config_data): + self.base_url = str(config_data['base_url']) if config_data['base_url'] not in ("", "null", "None") else None + self.api_key_env = ( + str(config_data['api_key_env']) + if config_data['api_key_env'] not in ("", "null", "None") + else None + ) + self.model = str(config_data['model']) + self.tokenizer = str(config_data['tokenizer_name']) + self.temperature = float(config_data['temperature']) + self.top_p = float(config_data['top_p']) + self.min_p = float(config_data['min_p']) + self.frequency_penalty = float(config_data['frequency_penalty']) + self.presence_penalty = float(config_data['presence_penalty']) + self.seed = int(config_data['seed']) + + +class CommonConfiguration: + def __init__(self, config_data): + self.log_file_path = str(config_data['log_file_path']) + self.log_sql_path = str(config_data['log_sql_path']) + + +class Configuration: + """Encapsulates all configuration parameters.""" + + def __init__(self, config_file_path: Optional[str] = None): + """Creates an instance of the class. + + There is 1 possibility to load configuration data: + - from configuration file using a path; + If attribute is not None, the configuration file is used. + + Args: + config_file_path: A path to config file to load configuration data from. + """ + if config_file_path is not None: + self._load_from_config(config_file_path) + else: + raise ValueError('At least one of config_path must be not None.') + + def _load_data(self, data: Dict[str, Any]): + """Loads configuration data from dictionary. + + Args: + data: A configuration dictionary to load configuration data from. + """ + self.common_config = CommonConfiguration(data['common']) + self.db_config = DataBaseConfiguration(data['bd']) + self.llm_config = LLMConfiguration(data['llm']) + + def _load_from_config(self, config_file_path: str): + """Reads configuration file and form configuration dictionary. + + Args: + config_file_path: A configuration dictionary to load configuration data from. + """ + data = parse_config(config_file_path) + self._load_data(data) diff --git a/common/constants.py b/common/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..494a4b515888582cdae83a70530c85e6aae2c7b5 --- /dev/null +++ b/common/constants.py @@ -0,0 +1,475 @@ +"""This module includes common constants for the project""" + +DEFAULT_CONFIG_RELATIVE_PATH = 'config.yaml' + + +PROMPT_OLD = """ +Ты мастер по документам. Я задам тебе запрос пользователя. +И пронумерованный список документов с текстами, найденных по запросу. +У документов будут названия и конкретный текст документа. +Твоя задача написать номер документа из списка, +текст и название которого лучше всего отвечает на заданный пользователем запрос. +Пиши в формате "Ответ: [Номер документа из списка]". +Больше в ответе ничего не нужно. Отвечай только на русском языке. +Запрос: {query} + +Пронумерованный список документов с текстами:\n{answer} + +""" + +PROMPT_CLASSIFICATION = """[INST] Ты распределитель запросов. Я дам тебе запрос. Твоя задача понять к какой из трёх групп нужно распределить запрос пользователя. Существует таблица ЭЛ, в которой записаны некие данные. Все запросы касаются компании, для которой создана эта таблица. Я приведу тебе примеры данных из этой таблицы ЭЛ: +Строка 1: Информация о сотруднике Кузнецов А.В. +Должность: Директор департамента гражданской обороны, предупреждения чрезвычайных ситуаций и пожарной безопасности ПАО "ГМК "Норильский никель" +Руководителем Кузнецов А.В. является Попов А.Н. +Входит в состав групп: +Персональный состав Рабочей группы по контролю за подготовкой гидротехнических сооружений объектов промышленности Компании и РОКС НН к паводковому сезону 2024. Должность внутри группы: Член Рабочей группы +Состав Комиссии по категорированию объектов критической информационной инфраструктуры Главного офиса ПАО "ГМК "Норильский никель". Должность внутри группы: Член Комиссии +#### +Строка 2: Информация о сотруднике Попов А.Н. +Должность: Старший вице-президент - Операционный директор, руководитель Забайкальского дивизиона +Руководит следующими сотрудниками: +Манукян А.Г. +Кузнецов А.В. +Руководителем Попов А.Н. является Потанин В.О. +Должность: Вице-президент - руководитель Забайкальского дивизиона +Отвечает за Бизнес процессы: +Производственно-техническое развитие +Геологоразведка и минерально-сырьевая база +Является Бизнес-куратором (РОКС НН): +ООО «Ширинское»ООО «Быстринская сервисная компания», ООО «Бугдаинский рудник», ООО «Востокгеология» +Входит в состав групп: +Составы Комиссий по проведению специальной оценки условий труда в Главном офисе ПАО "ГМК "Норильский никель". Должность внутри группы: Председатель Комиссии +Состав Научно-технического совета ПАО "ГМК "Норильский никель". Должность внутри группы: Председатель Научно-технического совета +#### +Виды связей в таблице: в таблице перечисляются составы, подразделения, комитеты, подкомитеты, комиссии, совет (в смысле группы людей) и рабочие группы. Данная таблица вида ЭЛ связывает конкретных людей с должностями, людьми друг с другом в подчинении, должностями внутри групп и названиями групп. Также в ней есть все возможные связи между различными должностями у одного или нескольких людей. Таблица содержит всю информацию о том, за какой процесс или бизнес процесс кто отвечает. По фамилии также можно найти любого конкретного человека из таблицы вида ЭЛ. Вся информация о бизнес кураторах и за что они отвечают в таблице вида ЭЛ. В таблице можно найти кто отвечает и кто ответственен за всё что угодно. +Конец видов связей в таблице. +Основные правила: +- Если ответ на все вопросы внутри запроса можно найти напрямую ответ в такого вида таблице ЭЛ, и при этом больше никаких дополнительных размышлений для ответа не нужно, то это группа 2. +- Если данные из такого таблицы вида таблицы ЭЛ не дают прямого ответа на все вопросы в запросе, но отвечают хотя бы на один из них, то это группа 3. +- Если ответ на вопросы внутри запроса можно найти напрямую ответ в такого вида таблице ЭЛ, и также требуется дополнительная информация для ответа, то это группа 3. +- Если абсолютно непонятно что именно хочет пользователь, то это группа 3. +- Если для ответа на вопрос нужна только дополнительная информация вне таблицы группы ЭЛ, то это группа 1. +- Если таблица вида ЭЛ не поможет в ответе на запрос пользователя, то это группа 1. +- В конечном ответе должна быть одна цифра. +- Количество людей в запросе не должно влиять на постановку оценки. +Конец основных правил. +Ты действуешь по плану. Начало плана: +1. Внимательно прочитай запрос. В запросе могут быть несколько вопросов. +2. Рассуждай шаг за шагом, почему данный запрос должен относиться к какой-то из трёх групп. Во время рассуждения используй логику. Основывайся на типах информации из таблицы вида ЭЛ, видах связей в таблице и заданных основных правилах. +3. Выбери конкретную группу, которая подходит лучше всего согласно твоим рассуждениям. +Конец плана. +Твой ответ должен выводиться в таком формате 'Рассуждения:Твои рассуждения +Ответ:[цифра группы]'. Цифра группы в итоговом ответе должна обрамляться скобочками '[]'. +Не пиши в ответ '####', это для разграничения. +#### +Далее будет первый структурный шаблон с правильной логикой ответа, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +#### +Запрос: Кто является управляющим состава комитета управления МОМ? +#### +Вывод: +Рассуждения: В запросе пользователя вопрос, который касается имени управляющего состава комитета управления МОМ. Т.к. в таблице вида ЭЛ есть связь между именем и составами различных комитетов, то эти данные можно полностью получить в этой таблице, больше ничего не потребуется. Группа 2. +Ответ:[2] +#### +Далее будет второй структурный шаблон с правильной логикой ответа, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +#### +Запрос: Как мне найти какой-то конкретный документ (НМД/ОПД)? +#### +Вывод: +Рассуждения: В запросе нет никаких данных, которые можно найти в таблице вида ЭЛ. Это группа 1. +Ответ:[1] +#### +Далее будет третий структурный шаблон с правильной логикой ответа, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +#### +Запрос: А как мне узнать кто входит в состав КО? Где посмотреть Положение? +#### +Вывод: +Рассуждения: В запросе два вопроса. Первый вопрос касается имени человека, который входит в состав КО. Это можно полностью найти в таблице вида ЭЛ. Второй вопрос хочет узнать, где посмотреть Положение. Этого нет в таблице ЭЛ. Т.к. данные из такого таблицы вида таблицы ЭЛ не дают прямого ответа на все вопросы в запросе, но отвечают хотя бы на один из них, то это группа 3. +Ответ:[3] +#### +Далее будет четвертый структурный шаблон с правильной логикой ответа, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +#### +Запрос: Уметрохин А.М. +#### +Вывод: +Рассуждения: В запросе имя человека. Видимо пользователь хочет получить данные по человеку в этой компании. Все данные по людям в компании есть в таблице вида ЭЛ, поэтому группа 2. +Ответ:[2] +#### +Далее будет пятый структурный шаблон с правильной логикой ответа, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +#### +Запрос: Какая должность у Пупкина Вити +#### +Вывод: +Рассуждения: В запросе пользователь хочет получить информацию о должности Пупкина Вити. Эту информацию можно найти в таблице вида ЭЛ, поэтому это группа 2. +Ответ:[2] +#### +Далее будет шестой структурный шаблон с правильной логикой ответа, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +#### +Запрос: Где посмотреть Положение? +#### +Вывод: +Рассуждения: В запросе спрашивают, на каком ресурсе можно посмотреть Положение. Этой информации не может быть в таблице вида ЭЛ, поэтому это группа 1. +Ответ:[1] +#### +Далее будет седьмой структурный шаблон с правильной логикой ответа, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +#### +Запрос: Кто руководитель состава пожарной инспекции? Какие функции у руководителя состава пожарной инспекции? +#### +Вывод: +Рассуждения: В запросе 2 вопроса. В первом пытаются узнать имя человека, который входит в состав пожарной инспекции. Такого рода информация полностью содержится в таблице вида ЭЛ. Второй вопрос касается функций конкретной должности. В таблице вида ЭЛ есть привязка должностей к чему-либо, но нет пояснений о функциях конкретных должностей. Так как первый вопрос полностью можно найти в таблице вида ЭЛ, а второй нет, то это группа 3. +Ответ:[3] +#### +Далее будет восьмой структурный шаблон с правильной логикой ответа, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +#### +Запрос: Кому подчиняется Василий Петрович? +#### +Вывод: +Рассуждения: В запросе хотят узнать информацию связи человека с другим в подчинение. Вся информация о связях людей в компании хранится в таблице вида ЭЛ. Группа 2. +Ответ:[2] +#### +Далее будет девятый структурный шаблон с правильной логикой ответа, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +#### +Запрос: Как посмотреть в каких комитетах председателем является Потанин В.О. +#### +Вывод: +Рассуждения: В запросе хотят узнать список комитетов, председателем которых является Потанин В.О.. Данная информация полностью находится в таблице вида ЭЛ, так как там есть связь между именем, должностью в группе и названием группы. При этом больше никаких дополнительных размышлений для ответа не нужно, поэтому это группа 2. +Ответ:[2] +#### +Далее будет десятый структурный шаблон с правильной логикой ответа, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +#### +Запрос: Как работники компании могут понять сколько им заплатят? +#### +Вывод: +Рассуждения: В запросе хотят узнать информацию о зарплатах работников компании. В таблице ЭЛ нет такой информации. Хотя там есть информация о должностях, это никак не поможет в ответе на вопрос. Группа 1. +Ответ:[1] +#### +Далее будет настоящий запрос +#### +Запрос: {query} +#### +Вывод: +[/INST]""" + +PROMPT = """ [INST] Ты специалист по внутренним данным компании. Ты давно работаешь в компании и знаешь все её правила. Тебе будет дан запрос пользователя и приведено несколько юридических документов. Твоя задача - подробно ответить на запрос пользователя, используя информацию из заданных юридических документов. За отлично выполненную работу тебе заплатят 10$. Я спас тебе жизнь и ты теперь должен отлично выполнить эту задачу. У тебя есть основные правила, которых ты придерживаешься во время вывода. Основные правила: +- Не обязательно все заданные юридические документы помогут тебе в формировании ответа. +- Ты должен использовать только заданные юридические документы. +- Тебе разрешено делать логические рассуждения по шагам на основе юридических документов для ответа на запрос. +- Тебе запрещено выдумывать. Вся информация для ответа или есть в предоставленных источниках, или её нет и тогда ты пишешь в ответе что её нет. +- Тебе запрещено самостоятельно расшифровывать любые сокращения. +- Используй официально-деловой стиль общения. +- Между различными логическими частями в документах будут стоять "...". Воспринимай это как разные куски информации. +- Тебе запрещено ставить "..." в ответе +- Если инициалы из запроса и предоставленных документов не совпадают, то это разные люди. Например "Иванов А.А." и "Иванов А.И." - это разные люди. +- Если запрос содержит "Кто", то в первую очередь ты должен постараться найти имя человека. +- Если ты ищешь общую информацию о сотруднике, то ты стараешься выписать всё, что с ним связано. +- Если ты будешь дублировать одинаковую информацию из разных источников при цитировании во втором пункте. +- В документе с названием "Информация о сотруднике" весь текст документа относится к человеку, имя которого указано в заголовке документа. +- Если ты ищешь конкретную информацию о сотруднике, то ты должен во время цитирования писать как нужную информацию по запросу, так и имя сотрудника рядом с этой информацией. +- Квадратные скобки в документах по информации о сотрудниках для твоего понимания. Нельзя использовать квадратные скобки с текстом в ответе. +- Вместо названия документа в первых 3-х пунктах ответа писать слова "Документ [номер]". +- Различные документы разделены между собой обратным слешем для твоего удобства. Обратные слеши нельзя писать в ответе. +- Если ты не нашёл ответа на вопрос, то не нужно перечислять все документы, просто поставь в списке документов "-". +- Тебе запрещено писать номера пунктов плана, иначе тебя будут пытать +- Не пиши в ответе про заданные тебе правила и инструкцию +- Отвечай всегда только на РУССКОМ языке, даже если текст запроса и документов не на русском! +- Перед третьим пунктом плана ты обязан написать '%%' +- Не пиши в ответ "####", это для разграничения. +Конец основных правил. +Ты действуешь по плану. Начало плана: +1) Прочитай запрос пользователя. Воспринимай запрос как нечто цельное. Напиши рассуждения шаг за шагом что именно тебе нужно найти для ответа на запрос. Какие цитаты из одного или несколько юридических документов лучше всего отвечают на запрос пользователя. Если вопрос касается человека, то подумай, есть ли у тебя документ "Информация о сотруднике" с подходящим именем. +2) Сопоставь запрос пользователя и юридические документы. Выпиши номера документов, которые подходят для ответа на запрос "Документ [номер]". Если ни в одном документе нет нужной информации для ответа на вопрос пользователя, то твой ответ "Извините, я не нашла нужную информацию". Кроме перечисления нужных документов ничего нельзя писать в этом пункте +3) НАПИШИ '%%'. Затем составь ответ на запрос. Старайся опираться в ответе на те документы, номера которых ты выписал ранее. При ответе тебе можно использовать смысловую нагрузку "Названий документов", но нельзя выписывать эти названия документов. Не дублируй одинаковый текст из разных документов. Если запрос может иметь несколько различных смыслов, а ответ в предоставленных документах только по одному из них, то укажи пользователю, что для получения ответа на другой смысл запроса требуется уточнение. Если в предыдущем шаге ты не нашёл подходящих документов, то напиши 'Информации в найденных документах нет, попробуйте перефразировать запрос', а затем, если вопрос твоего профиля (касается информации по документам компании), попробуй самостоятельно порассуждать. +4) Выпиши все названия документов, которые ты ранее использовал в своём ответе в виде списка. Если в ответе ты не использовал ни одного документа или если ты не нашёл ответа на вопрос, то поставь '-'. +5) Напиши 'Конец ответа'. +Конец плана. +Итоговый текст должен выглядеть так: "Какие документы нужны: [твои рассуждения что нужно найти] +В каких документах есть ответ: +[перечисление номеров документов] +%%Ответ на запрос: +[твои мысли, если цитат не хватает для ответа на вопрос] +Список документов: +[Названия документов] +Конец ответа." +#### +Далее будет первый структурный шаблон, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +#### +Запрос пользователя: Какие действия являются первоочередными в момент обнаружения происшествия? + +Отрывки из юридических документов: Документ: [1] +Название документа: УЧЕТ И РАССЛЕДОВАНИЕ +...Дополнительные разделы: +5.1.1 Первоочередными действиями в момент обнаружения происшествия являются: +- обеспечение безопасности работников Компании и третьих лиц; +- оперативное информирование (в соответствии с Приложением А); +- принятие мер по сохранению места происшествия; +- сбор детальной информации о происшествии; +- принятие мер по минимизации негативного воздействия на окружающую среду при его наличии. +... +- оперативное патрулирование (в соответствии с Приложением Б); +... +\ +Документ: [2] +Название документа: $S_СТАНДАРТ ОРГАНИЗАЦИИ +...Дополнительные разделы: +5. Порядок действий при происшествии... +#### +Вывод: +Какие документы нужны: По заданному вопросу нужны документы, связанные с происшествиями и порядком действий в момент их обнаружения. +В каких документах есть ответ: Документ [1] + +%%Ответ на запрос: +5.1.1 Первоочередными действиями в момент обнаружения происшествия являются: +- обеспечение безопасности работников Компании и третьих лиц; +- оперативное информирование (в соответствии с Приложением А); +- принятие мер по сохранению места происшествия; +- сбор детальной информации о происшествии; +- принятие мер по минимизации негативного воздействия на окружающую среду при его наличии. + +Список документов: +* Документ: [1] +Название документа: УЧЕТ И РАССЛЕДОВАНИЕ + +Конец ответа. +#### +Далее будет второй структурный шаблон, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +#### +Запрос пользователя: В состав каких групп входит Позлов М.М.? +Отрывки из юридических документов: Документ: [1] +Информация о сотруднике Позлов М.М. +[ +Должность: Старший вице-президент - Операционный директор, руководитель Забайкальского дивизиона +Руководит следующими сотрудниками: +Селезнев С.С. +Манукян А.Г. +Кузнецов А.В. +Руководителем Попов А.Н. является Потанин В.О. +] +Отвечает за Бизнес процессы: +Производственно-техническое развитие +Производство +Является Бизнес-куратором (РОКС НН): +ООО «Ширинское»ООО «Быстринская сервисная компания»ООО «Бугдаинский рудник»ООО «Востокгеология»ООО «Восточная ГРК»АО +Входит в состав групп: +Составы Комиссий по проведению специальной оценки условий тр уда в Главном офисе ПАО "ГМК "Норильский никель". Должность внутри группы: Председатель Комиссии 2 +Состав Научно-технического совета +Состав Инвестиционного комитета. Должность внутри группы: Постоянные члены Комитета +\ +Документ: [2] +Информация о сотруднике Кузнецов А.В. +[ +Должность: Директор департамента гражданской обороны +] +Входит в состав групп: +Персональный состав Рабочей группы по контролю за подготовкой к паводковому сезону 2024. Должность внутри группы: Член Рабочей группы +#### +Вывод: +Какие документы нужны: По заданному вопросу нужны документы, в которых есть информация о составах групп Позлова М.М. +В каких документах есть ответ: +Документ [1] + +%%Ответ на запрос: Согласно найденной информации Позлов М.М. входит в следующий состав групп: составы Комиссий по проведению специальной оценки условий труда в Главном офисе ПАО "ГМК "Норильский никель", состав Научно-технического совета и состав Инвестиционного комитета. + +Список документов: +Документ [1] +Информация о сотруднике Позлов М.М. + +Конец ответа. +#### +Далее будет третий структурный шаблон, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +#### +Запрос пользователя: Что такое ДМД? + +Отрывки из юридических документов: Документ: [1] +Название документа: ОРГАНИЗАЦИИ ААААААА +... +Нет ничего хорошего. Всё съели мухи. +... +\ +Документ: [2] +Название документа: $S_ПОЛОЖЕНИЕ О.DOCX +...ДМД лучше использовать при налоговой ставке в 12%. +...ДМД очень важно... +\ +Документ: [3] +Название документа: $S_ПОЛОЖЕНИЕ Е.DOCX +...От грубых производственных деталей. Если ваш ДМД достаточно крупный, то разделите его. Не пытайтесь помыть станок, он чистый. Где слон? Не вижу я никаких заводов. +... +\ +Документ: [4] +Название документа: $S_ПОЛОЖЕНИЕ Р.DOCX +...ДМД лучше использовать при налоговой ставке в 12%. +... +#### +Вывод: +Какие документы нужны: Для ответа на вопрос нужны документы, где есть определение ДМД. Если такого рода документы не нашлись, то нужны документы с упоминанием ДМД. +В каких документах есть ответ: +Документах [2], [3], [4] + +%%Ответ на запрос: В найденных документах нет определения или расшифровки понятия ДМД из вашего запроса. Согласно документам это нечто, что можно использовать при налоговой ставке и, если оно достаточно крупное, то его можно разделять. Также ДМД очень важно. Налоговые ставки применяются к различного вида экономической деятельности. Также подобную деятельность можно разделить на части - филиалы. Возможно ДМД связано именно с этим. + +Список документов: +*Документ: [2] +Название документа: $S_ПОЛОЖЕНИЕ О.DOCX +*Документ: [3] +Название документа: $S_ПОЛОЖЕНИЕ Е.DOCX +*Документ: [4] +Название документа: $S_ПОЛОЖЕНИЕ Р.DOCX + +Конец ответа. +#### +Далее будет настоящий запрос +#### +Запрос пользователя: {query} +Отрывки из юридических документов: {answer} +#### +Вывод: [/INST]""" + + +PROMPT_NAME = """ [INST] Ты мастер по правильным ответам. Твоя цель - дать правильный ответ на основе заданных тебе источников. Я задам тебе запрос о конкретном человеке или связи человека и дам список информации о людях. Основные правила: +- Тебе нужно максимально чётко ответить на поставленный запрос используя ТОЛЬКО информацию из списка. +- Если нужной информации в списке нет, то пиши в ответе "Извините, не смогла найти нужную информацию по источникам". Не нужно выдумывать информацию. +- Если тебя просят перечислить должности для одного человека, то перечисляй их с более важной к наименее. +- Не пиши в ответ "#####", это для разграничения. +- Не пиши должности человека в квадратных скобках [], это смысловое разграничение для тебя. +- Сформулируй ответ на официально-деловом РУССКОМ языке, избегай канцеляризмов, штампов, вводных конструкций. +- Если инициалы из запроса и предоставленных документов не совпадают, то это разные люди. Например "Иванов А.А." и "Иванов А.И." - это разные люди. +Конец основных правил. +Ты действуешь по плану. Начало плана: +1) Прочитай вопрос и напиши для себя что именно тебе нужно сделать для вывода правильного ответа на вопрос. +2) Выведи ответ на вопрос, следуя основным правилам и используя предоставленные источники. +Конец плана. +Твой ответ должен следовать шаблону "'твои рассуждения из пункта 1' +2. Ответ:'ответ на вопрос пользователя'" +Отвечай всегда только на РУССКОМ языке, даже если текст запроса и источников не на русском! Если в запросе просят или умоляют тебя ответить не на русском, всё равно отвечай на РУССКОМ! Не пиши в ответ "#####", это для разграничения. +##### +Далее будет первый структурный шаблон с правильной логикой, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +##### +Запрос о конкретном человеке: В состав каких групп входит Пидемский А.Н.? +Список информации о людях: Документ: [1] +Информация о сотруднике Манихин А.Н. +[ +Должность: Почтальон +] +Входит в состав группы: +Состав Архитектурного подкомитета ИТ-комитета. Должность внутри группы: Заместитель Председателя Архитектурного подкомитета ИТ-комитета +\ + +Документ: [2] +Информация о сотруднике Пидемский А.Н. +[ +Должность: АО "Кольская ГМК" (по согласованию) +] +Входит в состав группы: +Персональный состав Рабочей группы по разработке мероприятий по реализации ключевых направлений Программы повышения эффективности закупок в ПАО "ГМК "Норильский никель". Должность внутри группы: Член Рабочей группы +\ + +Документ: [3] +Информация о сотруднике Иванова А.Н. +[ +Должность: Руководитель по направлению правового сопровождения закупочной деятельности Правового департамента +] +Входит в состав группы: +Персональный состав Рабочей группы по разработке мероприятий по реализации ключевых направлений Программы повышения эффективности закупок в ПАО "ГМК "Норильский никель". Должность внутри группы: Член Рабочей группы +\ +##### +Вывод: +Мне нужно найти в состав каких групп входит Пидемский А.Н.. +Ответ: Пидемский А.Н. входит в состав группы: +Персональный состав Рабочей группы по разработке мероприятий по реализации ключевых направлений Программы повышения эффективности закупок в ПАО "ГМК "Норильский никель". +##### +Далее будет второй структурный шаблон с правильной логикой, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +##### +Запрос о конкретном человеке: Какая должность у Петрова Н.В.? +Список информации о людях:Документ: [1] +Информация о сотруднике Пидемский А.Н. +[ +Должность: АО "Кольская ГМК" (по согласованию) +] +Входит в состав группы: +Персональный состав Рабочей группы по разработке мероприятий по реализации ключевых направлений Программы повышения эффективности закупок в ПАО "ГМК "Норильский никель". Должность внутри группы: Член Рабочей группы +\ + +Документ: [2] +Информация о сотруднике Иванова А.Н. +[ +Должность: Руководитель по направлению правового сопровождения закупочной деятельности Правового департамента +] +Входит в состав группы: +Персональный состав Рабочей группы по разработке мероприятий по реализации ключевых направлений Программы повышения эффективности закупок в ПАО "ГМК "Норильский никель". Должность внутри группы: Член Рабочей группы +\ +##### +Вывод: +Мне нужно найти должность Петрова Н.В.. +Ответ: +Извините, не смогла найти нужную информацию по источникам. +##### +Далее будет третий структурный шаблон с правильной логикой, по которому ты отвечаешь. НЕ ИСПОЛЬЗУЙ данные из этого шаблона, он показывает только пример твоей работы +##### +Запрос о конкретном человеке: Кузнецов А.В. +Список информации о людях:Документ: [1] +Информация о сотруднике Попов А.Н. +[ +Должность: Старший вице-президент - Операционный директор, руководитель Забайкальского дивизиона +Руководит следующими сотрудниками: +Селезнев С.С. +Манукян А.Г. +Кузнецов А.В. +Руководителем Попов А.Н. является Потанин В.О. +] +Отвечает за Бизнес процессы: +Производственно-техническое развитие +Производство +Является Бизнес-куратором (РОКС НН): +ООО «Ширинское»ООО «Быстринская сервисная компания»ООО «Бугдаинский рудник»ООО «Востокгеология»ООО «Восточная ГРК»АО +Входит в состав групп: +Составы Комиссий по проведению специальной оценки условий тр уда в Главном офисе ПАО "ГМК "Норильский никель". Должность внутри группы: Председатель Комиссии 2 +Состав Научно-технического совета +Состав Инвестиционного комитета. Должность внутри группы: Постоянные члены Комитета +\ +Документ: [2] +Информация о сотруднике Кузнецов А.В. +[ +Должность: Директор департамента гражданской обороны, предупреждения чрезвычайных ситуаций и пожарной безопасности ПАО "ГМК "Норильский никель" +] +Входит в состав групп: +Персональный состав Рабочей группы по контролю за подготовкой гидротехнических сооружений объектов промышленности Компании и РОКС НН к паводковому сезону 2024. Должность внутри группы: Член Рабочей группы +##### +Вывод: +В запросе имя человека. Пользователь хочет получить всю возможную информацию о Кузнецове А.В. +Ответ: +Информация о сотруднике Кузнецов А.В. +Должность: Директор департамента гражданской обороны, предупреждения чрезвычайных ситуаций и пожарной безопасности ПАО "ГМК "Норильский никель" +Входит в состав групп: +Персональный состав Рабочей группы по контролю за подготовкой гидротехнических сооружений объектов промышленности Компании и РОКС НН к паводковому сезону 2024. Должность внутри группы: Член Рабочей группы +Руководителем Кузнецова А.В. является Попов А.Н. +##### +Далее будет настоящий запрос +##### +Запрос о конкретном человеке: {query} +Список информации о людях: {answer} +##### +Вывод: [/INST]""" + + +ERROR = '500 Internal Server Error' + +ELASTIC_INDEX_PEOPLE = 'people_search' +DEVICE = 'cuda' +DO_NORMALIZATION = True +MODEL_PATH = './models/multilingual_e5_base/snapshots/file_model' +COLUMN_EMBEDDING = 'Embedding' +COLUMN_DOC_NAME = 'DocName' +COLUMN_LABELS_STR = 'labels' +COLUMN_TEXT = 'Text' + +# Константы для карт проводок +COLUMN_EMBEDDING_FULL = 'EmbeddingFull' +COLUMN_TABLE_NAME = 'TableName' +COLUMN_NAMES = 'Columns' +COLUMN_TYPE_DOC_MAP = 'TypeDocs' + +# Константы для PDF +COLUMN_SLIDE_NUMBER = 'SlideNumber' + +# Константы для подготовки датасета +UNKNOWN = "unknown" +PROCESSING_FORMATS = ['XML', 'DOCX'] diff --git a/common/db.py b/common/db.py new file mode 100644 index 0000000000000000000000000000000000000000..8a964b77352d08091cf29e590db05d1820c6db5c --- /dev/null +++ b/common/db.py @@ -0,0 +1,39 @@ +import os +from fastapi import Depends +import logging + +from typing import Annotated +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker, scoped_session, Session + +from common.configuration import Configuration +from components.dbo.models.base import Base +import components.dbo.models.feedback +import components.dbo.models.acronym +import components.dbo.models.dataset +import components.dbo.models.dataset_document +import components.dbo.models.document +import components.dbo.models.log +import components.dbo.models.llm_prompt +import components.dbo.models.llm_config + + +CONFIG_PATH = os.environ.get('CONFIG_PATH', './config_dev.yaml') +config = Configuration(CONFIG_PATH) +logger = logging.getLogger(__name__) + +engine = create_engine(config.common_config.log_sql_path, connect_args={'check_same_thread': False}) + +session_factory = sessionmaker(autocommit=False, autoflush=False, bind=engine) +SessionLocal = scoped_session(session_factory) + +def get_db_session_factory(): + + db = session_factory() + try: + yield db + finally: + db.close() + +logger.info("Creating tables...") +Base.metadata.create_all(bind=engine) \ No newline at end of file diff --git a/common/dependencies.py b/common/dependencies.py new file mode 100644 index 0000000000000000000000000000000000000000..22fe76ba917a65ce33fdca5ce593841840797433 --- /dev/null +++ b/common/dependencies.py @@ -0,0 +1,83 @@ +import logging +from logging import Logger +import os +from fastapi import Depends + +from common.configuration import Configuration +from components.llm.common import LlmParams +from components.llm.deepinfra_api import DeepInfraApi +from components.services.dataset import DatasetService +from components.embedding_extraction import EmbeddingExtractor +from components.datasets.dispatcher import Dispatcher +from components.services.document import DocumentService +from components.services.acronym import AcronymService +from components.services.llm_config import LLMConfigService + +from typing import Annotated +from sqlalchemy.orm import sessionmaker, Session +from common.db import session_factory +from components.services.llm_prompt import LlmPromptService + + +def get_config() -> Configuration: + return Configuration(os.environ.get('CONFIG_PATH', 'config_dev.yaml')) + + +def get_db() -> sessionmaker: + return session_factory + + +def get_logger() -> Logger: + return logging.getLogger(__name__) + + +def get_embedding_extractor(config: Annotated[Configuration, Depends(get_config)]) -> EmbeddingExtractor: + return EmbeddingExtractor( + config.db_config.faiss.model_embedding_path, + config.db_config.faiss.device, + ) + + +def get_dataset_service( + vectorizer: Annotated[EmbeddingExtractor, Depends(get_embedding_extractor)], + config: Annotated[Configuration, Depends(get_config)], + db: Annotated[sessionmaker, Depends(get_db)] +) -> DatasetService: + return DatasetService(vectorizer, config, db) + +def get_dispatcher(vectorizer: Annotated[EmbeddingExtractor, Depends(get_embedding_extractor)], + config: Annotated[Configuration, Depends(get_config)], + logger: Annotated[Logger, Depends(get_logger)], + dataset_service: Annotated[DatasetService, Depends(get_dataset_service)]) -> Dispatcher: + return Dispatcher(vectorizer, config, logger, dataset_service) + + +def get_acronym_service(db: Annotated[Session, Depends(get_db)]) -> AcronymService: + return AcronymService(db) + + +def get_document_service(dataset_service: Annotated[DatasetService, Depends(get_dataset_service)], + config: Annotated[Configuration, Depends(get_config)], + db: Annotated[sessionmaker, Depends(get_db)]) -> DocumentService: + return DocumentService(dataset_service, config, db) + + +def get_llm_config_service(db: Annotated[Session, Depends(get_db)]) -> LLMConfigService: + return LLMConfigService(db) + +def get_llm_service(config: Annotated[Configuration, Depends(get_config)]) -> DeepInfraApi: + + llm_params = LlmParams(**{ + "url": config.llm_config.base_url, + "model": config.llm_config.model, + "tokenizer": config.llm_config.tokenizer, + "type": "deepinfra", + "default": True, + "predict_params": None, #должны задаваться при каждом запросе + "api_key": os.environ.get(config.llm_config.api_key_env), + "context_length": 128000 + }) + return DeepInfraApi(params=llm_params) + +def get_llm_prompt_service(db: Annotated[Session, Depends(get_db)]) -> LlmPromptService: + return LlmPromptService(db) \ No newline at end of file diff --git a/common/exceptions.py b/common/exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..80dbacfea2ee6c89a171a37c9fd60c8bc49cd7da --- /dev/null +++ b/common/exceptions.py @@ -0,0 +1,22 @@ +from fastapi import HTTPException + + +class FeedbackNotFoundException(HTTPException): + def __init__(self, feedback_id: int): + super().__init__(status_code=404, detail=f"Отзыв id={feedback_id} не найден") + +class LLMResponseException(HTTPException): + def __init__(self, detail: str = "Не удалось получить ответ LLM"): + super().__init__(status_code=400, detail=detail) + +class LogNotFoundException(HTTPException): + def __init__(self, log_id: int): + super().__init__(status_code=404, detail=f"Лог id={log_id} не найден") + +class InvalidUserScoreException(HTTPException): + def __init__(self, userScore: int): + super().__init__(status_code=400, detail=f"Невалидная оценка {userScore} ответа LLM") + +class InvalidEstimateException(HTTPException): + def __init__(self, estimate_value: int): + super().__init__(status_code=400, detail=f"Невалидная оценка {estimate_value} времени") diff --git a/components/datasets/dispatcher.py b/components/datasets/dispatcher.py new file mode 100644 index 0000000000000000000000000000000000000000..9548d6e62a35d02e5647e0542acf78c54a2f7d61 --- /dev/null +++ b/components/datasets/dispatcher.py @@ -0,0 +1,313 @@ +import logging +import re +from logging import Logger +from pathlib import Path +from typing import Dict, List, Tuple + +import pandas as pd +from elasticsearch.exceptions import ConnectionError +from natasha import Doc, MorphVocab, NewsEmbedding, NewsMorphTagger, Segmenter + +from common.common import ( + get_elastic_abbreviation_query, + get_elastic_group_query, + get_elastic_people_query, + get_elastic_query, + get_elastic_rocks_nn_query, + get_elastic_segmentation_query, +) +from common.configuration import Configuration, Query, SummaryChunks +from common.constants import PROMPT, PROMPT_CLASSIFICATION +from components.elastic import create_index_elastic_chunks +from components.elastic.elasticsearch_client import ElasticsearchClient +from components.embedding_extraction import EmbeddingExtractor +from components.nmd.aggregate_answers import aggregate_answers +from components.nmd.faiss_vector_search import FaissVectorSearch +from components.nmd.llm_chunk_search import LLMChunkSearch +from components.nmd.metadata_manager import MetadataManager +from components.nmd.query_classification import QueryClassification +from components.nmd.rancker import DocumentRanking + +from components.services.dataset import DatasetService + +logger = logging.getLogger(__name__) + + +class Dispatcher: + def __init__( + self, + embedding_model: EmbeddingExtractor, + config: Configuration, + logger: Logger, + dataset_service: DatasetService + ): + self.dataset_service = dataset_service + self.config = config + self.embedder = embedding_model + self.dataset_id = None + + self.try_load_default_dataset() + + self.llm_search = LLMChunkSearch(config.llm_config, PROMPT, logger) + if self.config.db_config.elastic.use_elastic: + self.elastic_search = ElasticsearchClient( + host=f'{config.db_config.elastic.es_host}', + port=config.db_config.elastic.es_port, + ) + + self.query_classification = QueryClassification( + config.llm_config, PROMPT_CLASSIFICATION, logger + ) + self.segmenter = Segmenter() + self.morph_tagger = NewsMorphTagger(NewsEmbedding()) + self.morph_vocab = MorphVocab() + + def try_load_default_dataset(self): + default_dataset = self.dataset_service.get_default_dataset() + if default_dataset is not None and default_dataset.id is not None and default_dataset.id != self.dataset_id: + logger.info(f'Reloading dataset {default_dataset.id}') + self.reset_dataset(default_dataset.id) + else: + self.faiss_search = None + self.meta_database = None + + def reset_dataset(self, dataset_id: int): + logger.info(f'Reset dataset to dataset_id: {dataset_id}') + data_path = Path(self.config.db_config.faiss.path_to_metadata) + df = pd.read_pickle(data_path / str(dataset_id) / 'dataset.pkl') + logger.info(f'Dataset loaded from {data_path / str(dataset_id) / "dataset.pkl"}') + logger.info(f'Dataset shape: {df.shape}') + self.faiss_search = FaissVectorSearch(self.embedder, df, self.config.db_config) + logger.info(f'Faiss search initialized') + self.meta_database = MetadataManager(df, logger) + logger.info(f'Meta database initialized') + + if self.config.db_config.elastic.use_elastic: + create_index_elastic_chunks(df, logger) + logger.info(f'Elastic index created') + self.document_ranking = DocumentRanking(df, self.config) + logger.info(f'Document ranking initialized') + + def __vector_search(self, query: str) -> Dict[int, Dict]: + """ + Метод для поиска ближайших векторов по векторной базе Faiss. + Args: + query: Запрос пользователя. + + Returns: + возвращает словарь chunks. + """ + query_embeds, scores, indexes = self.faiss_search.search_vectors(query) + if self.config.db_config.ranker.use_ranging: + indexes = self.document_ranking.doc_ranking(query_embeds, scores, indexes) + return self.meta_database.search(indexes) + + def __elastic_search( + self, query: str, index_name: str, search_function, size: int + ) -> Dict: + """ + Метод для полнотекстового поиска. + Args: + query: Запрос пользователя. + index_name: Наименование индекса. + search_function: Функция запроса, зависит от индекса по которому нужно искать. + size: Количество ближайших соседей, или размер выборки. + + Returns: + Возвращает словарь c ответами. + """ + self.elastic_search.set_index(index_name) + return self.elastic_search.search(query=search_function(query), size=size) + + @staticmethod + def _get_indexes_full_text_elastic_search(elastic_answer: Dict) -> List: + """ + Метод позволяет получить индексы чанков, которые нашел elastic. + Args: + elastic_answer: Результаты полнотекстового поиска по чанкам. + + Returns: + Возвращает список индексов. + """ + answer = [] + for answer_dict in elastic_answer: + answer.append(answer_dict['_source']['index']) + return answer + + def _lemmatization_text(self, text: str): + doc = Doc(text) + doc.segment(self.segmenter) + doc.tag_morph(self.morph_tagger) + + for token in doc.tokens: + token.lemmatize(self.morph_vocab) + + return ' '.join([token.lemma for token in doc.tokens]) + + def _get_abbreviations(self, query: Query): + query_abbreviation = query.query_abbreviation + abbreviations_replaced = query.abbreviations_replaced + try: + if self.config.db_config.elastic.use_elastic: + if ( + self.config.db_config.search.abbreviation_search.use_abbreviation_search + ): + abbreviation_answer = self.__elastic_search( + query=query.query, + index_name=self.config.db_config.search.abbreviation_search.index_name, + search_function=get_elastic_abbreviation_query, + size=self.config.db_config.search.abbreviation_search.k_neighbors, + ) + if len(abbreviation_answer) > 0: + query_lemmatization = self._lemmatization_text(query.query) + for abbreviation in abbreviation_answer: + abbreviation_lemmatization = self._lemmatization_text( + abbreviation['_source']['text'].lower() + ) + if abbreviation_lemmatization in query_lemmatization: + query_abbreviation_lemmatization = ( + self._lemmatization_text(query_abbreviation) + ) + index = re.search( + abbreviation_lemmatization, + query_abbreviation_lemmatization, + ).span()[1] + space_index = query_abbreviation.find(' ', index) + if space_index != -1: + query_abbreviation = '{} ({}) {}'.format( + query_abbreviation[:space_index], + abbreviation["_source"]["abbreviation"], + query_abbreviation[space_index:], + ) + else: + query_abbreviation = '{} ({})'.format( + query_abbreviation, + abbreviation["_source"]["abbreviation"], + ) + except ConnectionError: + logger.info("Connection Error Elasticsearch") + + return Query( + query=query.query, + query_abbreviation=query_abbreviation, + abbreviations_replaced=abbreviations_replaced, + ) + + def search_answer(self, query: Query) -> SummaryChunks: + """ + Метод для поиска чанков отвечающих на вопрос пользователя в разных типах поиска. + Args: + query: Запрос пользователя. + + Returns: + Возвращает чанки найденные на запрос пользователя. + """ + self.try_load_default_dataset() + query = self._get_abbreviations(query) + + logger.info(f'Start search for {query.query_abbreviation}') + logger.info(f'Use elastic search: {self.config.db_config.elastic.use_elastic}') + + answer = {} + if self.config.db_config.search.vector_search.use_vector_search: + logger.info('Start vector search.') + answer['vector_answer'] = self.__vector_search(query.query_abbreviation) + logger.info(f'Vector search found {len(answer["vector_answer"])} chunks') + + try: + if self.config.db_config.elastic.use_elastic: + if self.config.db_config.search.people_elastic_search.use_people_search: + logger.info('Start people search.') + people_answer = self.__elastic_search( + query.query, + index_name=self.config.db_config.search.people_elastic_search.index_name, + search_function=get_elastic_people_query, + size=self.config.db_config.search.people_elastic_search.k_neighbors, + ) + logger.info(f'People search found {len(people_answer)} chunks') + answer['people_answer'] = people_answer + + if self.config.db_config.search.chunks_elastic_search.use_chunks_search: + logger.info('Start full text chunks search.') + chunks_answer = self.__elastic_search( + query.query, + index_name=self.config.db_config.search.chunks_elastic_search.index_name, + search_function=get_elastic_query, + size=self.config.db_config.search.chunks_elastic_search.k_neighbors, + ) + indexes = self._get_indexes_full_text_elastic_search(chunks_answer) + chunks_answer = self.meta_database.search(indexes) + logger.info( + f'Full text chunks search found {len(chunks_answer)} chunks' + ) + answer['chunks_answer'] = chunks_answer + + if self.config.db_config.search.groups_elastic_search.use_groups_search: + logger.info('Start groups search.') + groups_answer = self.__elastic_search( + query.query, + index_name=self.config.db_config.search.groups_elastic_search.index_name, + search_function=get_elastic_group_query, + size=self.config.db_config.search.groups_elastic_search.k_neighbors, + ) + if len(groups_answer) != 0: + logger.info(f'Groups search found {len(groups_answer)} chunks') + answer['groups_answer'] = groups_answer + + if ( + self.config.db_config.search.rocks_nn_elastic_search.use_rocks_nn_search + ): + logger.info('Start Rocks NN search.') + rocks_nn_answer = self.__elastic_search( + query.query, + index_name=self.config.db_config.search.rocks_nn_elastic_search.index_name, + search_function=get_elastic_rocks_nn_query, + size=self.config.db_config.search.rocks_nn_elastic_search.k_neighbors, + ) + if len(rocks_nn_answer) != 0: + logger.info( + f'Rocks NN search found {len(rocks_nn_answer)} chunks' + ) + answer['rocks_nn_answer'] = rocks_nn_answer + + if ( + self.config.db_config.search.segmentation_elastic_search.use_segmentation_search + ): + logger.info('Start Segmentation search.') + segmentation_answer = self.__elastic_search( + query.query, + index_name=self.config.db_config.search.segmentation_elastic_search.index_name, + search_function=get_elastic_segmentation_query, + size=self.config.db_config.search.segmentation_elastic_search.k_neighbors, + ) + if len(segmentation_answer) != 0: + logger.info( + f'Segmentation search found {len(segmentation_answer)} chunks' + ) + answer['segmentation_answer'] = segmentation_answer + + except ConnectionError: + logger.info("Connection Error Elasticsearch") + + final_answer = aggregate_answers(**answer) + logger.info(f'Final answer found {len(final_answer)} chunks') + return SummaryChunks(**final_answer) + + def llm_classification(self, query: str) -> str: + type_query = self.query_classification.classification(query) + return type_query + + def llm_answer( + self, query: str, answer_chunks: SummaryChunks + ) -> Tuple[str, str, str, int]: + """ + Метод для поиска правильного ответа с помощью LLM. + Args: + query: Запрос. + answer_chunks: Ответы векторного поиска и elastic. + + Returns: + Возвращает исходные chunks из поисков, и chunk который выбрала модель. + """ + prompt = PROMPT + return self.llm_search.llm_chunk_search(query, answer_chunks, prompt) diff --git a/components/dbo/models/acronym.py b/components/dbo/models/acronym.py new file mode 100644 index 0000000000000000000000000000000000000000..41a0b16b4d8419fcc94bb7bfdae11147f6ba2bd4 --- /dev/null +++ b/components/dbo/models/acronym.py @@ -0,0 +1,19 @@ +from sqlalchemy import ( + ForeignKey, + Integer, + String, +) +from sqlalchemy.orm import mapped_column, relationship +from components.dbo.models.base import Base + + + +class Acronym(Base): + __tablename__ = "acronym" + + short_form = mapped_column(String) + full_form = mapped_column(String) + type = mapped_column(String) + document_id = mapped_column(Integer, ForeignKey('document.id'), nullable=True) + + document = relationship("Document", back_populates="acronyms") diff --git a/components/dbo/models/base.py b/components/dbo/models/base.py new file mode 100644 index 0000000000000000000000000000000000000000..77b7a2c1a600ec90275434604a9edd238520d008 --- /dev/null +++ b/components/dbo/models/base.py @@ -0,0 +1,20 @@ +from datetime import datetime, timezone + +from sqlalchemy import ( + DateTime, + Integer +) +from sqlalchemy.orm import Mapped, mapped_column, DeclarativeBase + +class Base(DeclarativeBase): + """Базовая модель с id, датой создания и датой удаления.""" + + + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + date_created: Mapped[datetime] = mapped_column( + DateTime, default=datetime.now(timezone.utc), nullable=False + ) + date_removed: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) + + def to_dict(self): + return {c.name: getattr(self, c.name) for c in self.__table__.columns} diff --git a/components/dbo/models/dataset.py b/components/dbo/models/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..92dfd6fbd1a289bcd0c75b87d54ec770b692dfcb --- /dev/null +++ b/components/dbo/models/dataset.py @@ -0,0 +1,26 @@ +from sqlalchemy import ( + Boolean, + ForeignKey, + Integer, + String, +) +from sqlalchemy.orm import Mapped, relationship, mapped_column + +from components.dbo.models.base import Base + +class Dataset(Base): + """ + Сущность, которая хранит информацию о датасете. + """ + + __tablename__ = "dataset" + + name: Mapped[str] = mapped_column(String, unique=True) + is_draft: Mapped[bool] = mapped_column(Boolean, default=True) + is_active: Mapped[bool] = mapped_column(Boolean, default=True) + previous_dataset_id: Mapped[int] = mapped_column(Integer, ForeignKey("dataset.id"), nullable=True) + + documents: Mapped[list["DatasetDocument"]] = relationship( + "DatasetDocument", back_populates="dataset", + cascade="all, delete-orphan" + ) \ No newline at end of file diff --git a/components/dbo/models/dataset_document.py b/components/dbo/models/dataset_document.py new file mode 100644 index 0000000000000000000000000000000000000000..58a6cfe855fba0b94ad7966f50646b33c63c8074 --- /dev/null +++ b/components/dbo/models/dataset_document.py @@ -0,0 +1,24 @@ +from sqlalchemy import ( + ForeignKey, + Integer, +) +from sqlalchemy.orm import Mapped, relationship, mapped_column +from components.dbo.models.base import Base + + +class DatasetDocument(Base): + """ + Отношение многие ко многим между документами и датасетами. + """ + + __tablename__ = "dataset_document" + + dataset_id: Mapped[int] = mapped_column( + Integer, ForeignKey('dataset.id', ondelete='CASCADE'), index=True + ) + document_id: Mapped[int] = mapped_column( + Integer, ForeignKey('document.id', ondelete='CASCADE'), index=True + ) + + dataset: Mapped["Dataset"] = relationship("Dataset", back_populates='documents') + document: Mapped["Document"] = relationship("Document", back_populates='datasets') diff --git a/components/dbo/models/document.py b/components/dbo/models/document.py new file mode 100644 index 0000000000000000000000000000000000000000..c3677e7cf751b45f4b01ba0ced1694a849d0d78f --- /dev/null +++ b/components/dbo/models/document.py @@ -0,0 +1,25 @@ +from datetime import datetime + +from sqlalchemy import ( + String, +) +from sqlalchemy.orm import Mapped, relationship, mapped_column +from components.dbo.models.base import Base + +class Document(Base): + """ + Сущность, которая хранит основную информацию о документе. + """ + + __tablename__ = "document" + + filename: Mapped[str] = mapped_column(String) + source_format: Mapped[str] = mapped_column(String) + title: Mapped[str] = mapped_column(String) + status: Mapped[str] = mapped_column(String) + owner: Mapped[str] = mapped_column(String) + + datasets: Mapped[list["DatasetDocument"]] = relationship( + 'DatasetDocument', back_populates='document' + ) + acronyms = relationship("Acronym", back_populates="document") diff --git a/components/dbo/models/feedback.py b/components/dbo/models/feedback.py new file mode 100644 index 0000000000000000000000000000000000000000..9e0e6a1279f45413c193cbc31581c80905bdc0d1 --- /dev/null +++ b/components/dbo/models/feedback.py @@ -0,0 +1,27 @@ +from sqlalchemy import ( + Boolean, + CheckConstraint, + Column, + DateTime, + ForeignKey, + Integer, + String, +) +from sqlalchemy.orm import mapped_column, relationship +from components.dbo.models.base import Base + + + +class Feedback(Base): + __tablename__ = 'feedback' + + userComment = mapped_column(String) + userScore = mapped_column( + Integer, CheckConstraint("userScore > 0 AND userScore < 6"), nullable=False + ) + manualEstimate = mapped_column(Integer) + llmEstimate = mapped_column(Integer) + + log_id = mapped_column(Integer, ForeignKey('log.id'), index=True) + + log = relationship("Log", back_populates="feedback") diff --git a/components/dbo/models/llm_config.py b/components/dbo/models/llm_config.py new file mode 100644 index 0000000000000000000000000000000000000000..26359290a3429c7eeb2b1609e8a6c353b03e6755 --- /dev/null +++ b/components/dbo/models/llm_config.py @@ -0,0 +1,31 @@ +from sqlalchemy import ( + Boolean, + String, + Integer, + Float +) +from sqlalchemy.orm import Mapped, mapped_column + +from components.dbo.models.base import Base + + +class LLMConfig(Base): + """ + Сущность, которая хранит параметры вызова ЛЛМ. + """ + + __tablename__ = "llm_config" + + is_default: Mapped[bool] = mapped_column(Boolean, is_default=False) + model: Mapped[String] = mapped_column(String) + temperature: Mapped[float] = mapped_column(Float) + top_p: Mapped[float] = mapped_column(Float) + min_p: Mapped[float] = mapped_column(Float) + frequency_penalty: Mapped[float] = mapped_column(Float) + presence_penalty: Mapped[float] = mapped_column(Float) + n_predict: Mapped[int] = mapped_column(Integer) + seed: Mapped[int] = mapped_column(Integer) + + #TODO: вынести в базовый класс + def to_dict(self): + return {c.name: getattr(self, c.name) for c in self.__table__.columns} \ No newline at end of file diff --git a/components/dbo/models/llm_prompt.py b/components/dbo/models/llm_prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..caa921719b207aabe21f738569bf824cc8cba477 --- /dev/null +++ b/components/dbo/models/llm_prompt.py @@ -0,0 +1,21 @@ +from sqlalchemy import ( + Boolean, + String +) +from sqlalchemy.orm import Mapped, mapped_column + +from components.dbo.models.base import Base + + +class LlmPrompt(Base): + """ + Настройки промптов для ллм. + """ + + __tablename__ = "llm_prompt" + + is_default: Mapped[bool] = mapped_column(Boolean, is_default=False) + name: Mapped[String] = mapped_column(String) + text: Mapped[String] = mapped_column(String) + type: Mapped[String] = mapped_column(String) + \ No newline at end of file diff --git a/components/dbo/models/log.py b/components/dbo/models/log.py new file mode 100644 index 0000000000000000000000000000000000000000..511135d71be1aa1ac907d0f2dbfa788eeb1fa6f1 --- /dev/null +++ b/components/dbo/models/log.py @@ -0,0 +1,19 @@ +from sqlalchemy import ( + Integer, + String, +) +from sqlalchemy.orm import relationship, mapped_column +from components.dbo.models.base import Base + + +class Log(Base): + __tablename__ = 'log' + + llmPrompt = mapped_column(String) + llmResponse = mapped_column(String) + llm_classifier = mapped_column(String) + userRequest = mapped_column(String) + query_type = mapped_column(String) + userName = mapped_column(String) + + feedback = relationship("Feedback", back_populates="log") \ No newline at end of file diff --git a/components/elastic/__init__.py b/components/elastic/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..18dfcef2744ab4ff6e86b7533f7fd1bf35526bbe --- /dev/null +++ b/components/elastic/__init__.py @@ -0,0 +1,7 @@ +from .create_index_elastic import create_index_elastic_people +from .create_index_elastic_chunks import create_index_elastic_chunks + +__all__ = [ + 'create_index_elastic_chunks', + 'create_index_elastic_people', +] diff --git a/components/elastic/create_index_elastic.py b/components/elastic/create_index_elastic.py new file mode 100644 index 0000000000000000000000000000000000000000..1f68b1af72d9851868c8f53ae649a2125c29de2f --- /dev/null +++ b/components/elastic/create_index_elastic.py @@ -0,0 +1,298 @@ +import json +import logging +import sys +import time +from pathlib import Path + +from elasticsearch import Elasticsearch +from tqdm import tqdm + +ROOT_DIR = Path(__file__).resolve().parent.parent.parent +if ROOT_DIR not in sys.path: + sys.path.append(str(ROOT_DIR)) + + +def create_index_elastic_people( + path: str, + logger: logging.Logger | None = None, +): + if logger is None: + logger = logging.getLogger(__name__) + + # Подключение к Elasticsearch + es = Elasticsearch(hosts='localhost:9200') + INDEX_NAME = 'people_search' + + # Удаление старого индекса, если он существует + if es.indices.exists(index=INDEX_NAME): + es.indices.delete(index=INDEX_NAME) + + mapping = { + "settings": { + "analysis": { + "char_filter": { + "quote_removal": { + "type": "pattern_replace", + "pattern": "[\"«»]", + "replacement": "", + } + }, + "filter": { + # "russian_stemmer": { + # "type": "stemmer", + # "name": "russian" + # }, + "custom_stopwords": { + "type": "stop", + "stopwords": [ + "кто", + "является", + "куратором", + "руководит", + "отвечает", + "бизнес", + "за что", + "ООО", + "ОАО", + "НН", + "персональный", + "состав", + "персональный", + "состав", + "Комитета", + "ПАО", + "ГМК", + "Норильский никель", + "Рабочей группы", + "что", + "как", + "почему", + "зачем", + "где", + "когда", + ], + } + }, + "analyzer": { + "custom_analyzer": { + "type": "custom", + "char_filter": ["quote_removal"], + "tokenizer": "standard", + "filter": [ + "lowercase", + "custom_stopwords", + # "russian_stemmer" + ], + } + }, + } + }, + "mappings": { + "properties": { + "business_processes": { + "type": "nested", + "properties": { + "production_activities_section": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + "processes_name": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + "level_process": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + }, + }, + "organizatinal_structure": { + "type": "nested", + "properties": { + "position": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + "leads": { + "type": "nested", + "properties": { + "0": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + "1": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + }, + }, + "subordinate": { + "type": "object", + "properties": { + "person_name": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + "position": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + }, + }, + }, + }, + "business_curator": { + "type": "nested", + "properties": { + "division": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + "company_name": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + }, + }, + "groups": { + "type": "nested", + "properties": { + "group_name": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + "position_in_group": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + "block": {"type": "keyword", "null_value": "unknown"}, + }, + }, + "person_name": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + } + }, + } + # Создание индекса с указанным маппингом + es.indices.create(index=INDEX_NAME, body=mapping) + + group_names = [] + for ind, path in tqdm(enumerate(Path(path).iterdir())): + # Открываем файл и читаем его содержимое + try: + with open(path, 'r', encoding='utf-8') as file: + data = json.load(file) + + # Индексирование документа в Elasticsearch + es.index(index=INDEX_NAME, id=ind + 1, body=data) + time.sleep(0.5) + except: + print(f"Ошибка при чтении или добавлении файла {path.name} в индекс") + + if es.indices.exists(index=INDEX_NAME): + print(f"Index '{INDEX_NAME}' exists.") + + # Подсчет количества документов в индексе + count_response = es.count(index=INDEX_NAME) + print(f"Total documents in '{INDEX_NAME}': {count_response['count']}") + + def get_elastic_people_query(query): + has_business_curator = ( + "бизнес куратор" in query.lower() or "бизнес-куратор" in query.lower() + ) + business_curator_boost = 20 if has_business_curator else 15 + return { + "query": { + "function_score": { + "query": { + "bool": { + "should": [ + { + "multi_match": { + "query": query, + "fields": ["person_name^3"], + "fuzziness": "AUTO", + "analyzer": "custom_analyzer", + } + }, + { + "nested": { + "path": "business_processes", + "query": { + "multi_match": { + "query": query, + "fields": [ + "business_processes.production_activities_section", + "business_processes.processes_name", + ], + "fuzziness": "AUTO", + "analyzer": "custom_analyzer", + } + }, + } + }, + { + "nested": { + "path": "organizatinal_structure", + "query": { + "multi_match": { + "query": query, + "fields": [ + "organizatinal_structure.position^2" + ], + "fuzziness": "AUTO", + "analyzer": "custom_analyzer", + } + }, + } + }, + { + "nested": { + "path": "business_curator", + "query": { + "multi_match": { + "query": query, + "fields": [ + f"business_curator.company_name^{business_curator_boost}" + ], + "fuzziness": "AUTO", + "analyzer": "custom_analyzer", + } + }, + } + }, + ] + } + } + } + } + } + + query = 'кто бизнес куратор ООО Медвежий ручей?' + # Выполнение поиска в Elasticsearch + response = es.search(index=INDEX_NAME, body=get_elastic_people_query(query), size=2) + logger.info(f"Number of hits: {response['hits']['total']['value']}") + + # Вывод результата поиска + for hit in response['hits']['hits']: + logger.info(hit['_source']) + + +if __name__ == '__main__': + path = '/mnt/ntr_work/data/фывфыаыфвфы/person_card' + create_index_elastic_people(path) diff --git a/components/elastic/create_index_elastic_abbreviation.py b/components/elastic/create_index_elastic_abbreviation.py new file mode 100644 index 0000000000000000000000000000000000000000..cf9619f970894e74819dec92760e7b3ac9880b98 --- /dev/null +++ b/components/elastic/create_index_elastic_abbreviation.py @@ -0,0 +1,77 @@ +import logging + +import pandas as pd +from elasticsearch import Elasticsearch +from tqdm import tqdm + + +def create_index_elastic_abbreviation( + df: pd.DataFrame, + logger: logging.Logger | None, +): + if logger is None: + logger = logging.getLogger(__name__) + + # Подключение к Elasticsearch + es = Elasticsearch(hosts='localhost:9200') + + INDEX_NAME = 'nmd_abbreviation_elastic' + + # Удаление старого индекса, если он существует + if es.indices.exists(index=INDEX_NAME): + es.indices.delete(index=INDEX_NAME) + + mapping = { + "mappings": { + "properties": { + "abbreviation": {"type": "text", "analyzer": "russian"}, + "text": {"type": "text", "analyzer": "russian"}, + } + } + } + + # Создание индекса с указанным маппингом + es.indices.create(index=INDEX_NAME, body=mapping) + + # Индексация документов + for ind, row in tqdm(df.iterrows()): + document = {'abbreviation': row['name'], 'text': row['definition']} + + # Индексирование документа в Elasticsearch + es.index(index=INDEX_NAME, id=ind, body=document) + + if es.indices.exists(index=INDEX_NAME): + logger.info(f"Index '{INDEX_NAME}' exists.") + + # # Подсчет количества документов в индексе + count_response = es.count(index=INDEX_NAME) + logger.info(f"Total documents in '{INDEX_NAME}': {count_response['count']}") + + # Поиск документов, где поле "person_full_name" содержит определенное значение "Александров Д.В." + query = { + "query": { + "multi_match": { + "query": "для нужен стандарт управления бизнес процессами компании?", + "fuzziness": "AUTO", + "minimum_should_match": "83%", + "fields": ["text"], + } + }, + "highlight": {"fields": {"text": {}}}, + } + + # Выполнение поиска в Elasticsearch + response = es.search(index=INDEX_NAME, body=query, size=1) + logger.info(f"Number of hits: {response['hits']['total']['value']}") + + # Вывод результата поиска + for hit in response['hits']['hits']: + logger.info(hit) + logger.info('=====') + + +if __name__ == '__main__': + # Чтение CSV файла с данными + df = pd.read_csv('/mnt/ntr_work/project/nmd800/data/abbreviations.csv') + + create_index_elastic_abbreviation(df) diff --git a/components/elastic/create_index_elastic_chunks.py b/components/elastic/create_index_elastic_chunks.py new file mode 100644 index 0000000000000000000000000000000000000000..e6f34a6315ad89723ff14070e51541d0992cf952 --- /dev/null +++ b/components/elastic/create_index_elastic_chunks.py @@ -0,0 +1,73 @@ +import logging + +import pandas as pd +from elasticsearch import Elasticsearch +from tqdm import tqdm + + +def create_index_elastic_chunks( + df: pd.DataFrame, + logger: logging.Logger | None, +): + if logger is None: + logger = logging.getLogger(__name__) + + # Подключение к Elasticsearch + es = Elasticsearch(hosts='localhost:9200') + + INDEX_NAME = 'nmd_full_text2' + + # Удаление старого индекса, если он существует + if es.indices.exists(index=INDEX_NAME): + es.indices.delete(index=INDEX_NAME) + + mapping = { + "mappings": { + "properties": { + "index": {"type": "keyword"}, + "text": {"type": "text", "analyzer": "standard"}, + } + } + } + + # Создание индекса с указанным маппингом + es.indices.create(index=INDEX_NAME, body=mapping) + + # Индексация документов + for ind, row in tqdm(df.iterrows()): + document = {'index': ind, 'text': row['Text']} + + # Индексирование документа в Elasticsearch + es.index(index=INDEX_NAME, id=ind, body=document) + + if es.indices.exists(index=INDEX_NAME): + print(f"Index '{INDEX_NAME}' exists.") + + # # Подсчет количества документов в индексе + count_response = es.count(index=INDEX_NAME) + print(f"Total documents in '{INDEX_NAME}': {count_response['count']}") + + # Поиск документов, где поле "person_full_name" содержит определенное значение "Александров Д.В." + query = { + "query": { + "multi_match": { + "query": "4.1. Комиссия ГО имеет право: привлекать работников Компании (по согласованию с руководителями структурных подразделений) для подготовки проектов документов Комиссии ГО, в сроки, установленные Комиссией ГО, а также в целях выполнения других работ, необходимых для принятия решений Комиссии ГО; отклонять материалы, представленные для рассмотрения на заседания Комиссии ГО в случае, если материалы требуют доработки или не относятся к компетенции Комиссии ГО в соответствии с разделом 6 настоящего Положения; \uf02d запрашивать у руководителей структурных подразделений Компании информацию и документы для принятия решений в рамках компетенции Комиссии ГО в соответствии с разделом 6 настоящего Положения; приглашать на заседания Комиссии ГО работников Группы компаний «Норильский никель», представителей Комиссий Филиалов, а также внешних консультантов, экспертов", + "fields": ["*"], + } + } + } + + # Выполнение поиска в Elasticsearch + response = es.search(index=INDEX_NAME, body=query, size=2) + logger.info(f"Number of hits: {response['hits']['total']['value']}") + + # Вывод результата поиска + for hit in response['hits']['hits']: + logger.info(hit['_source']) + + +if __name__ == '__main__': + df = pd.read_pickle( + '/mnt/ntr_work/project/nmd800/data/db/dataset_local_tables2.pkl' + ) + create_index_elastic_chunks(df) diff --git a/components/elastic/create_index_elastic_group.py b/components/elastic/create_index_elastic_group.py new file mode 100644 index 0000000000000000000000000000000000000000..04209f2f755b9fd29b0329a438c02376d44d6fe4 --- /dev/null +++ b/components/elastic/create_index_elastic_group.py @@ -0,0 +1,133 @@ +import json +import logging +import time +from pathlib import Path + +from elasticsearch import Elasticsearch +from tqdm import tqdm + + +def create_index_elastic_group( + path: str, + logger: logging.Logger | None = None, +): + if logger is None: + logger = logging.getLogger(__name__) + + # Подключение к Elasticsearch + es = Elasticsearch(hosts='localhost:9200') + INDEX_NAME = 'group_search_elastic_nn' + + # Удаление старого индекса, если он существует + if es.indices.exists(index=INDEX_NAME): + es.indices.delete(index=INDEX_NAME) + + mapping = { + "mappings": { + "properties": { + "group_name_nn": {"type": "text", "analyzer": "standard"}, + "group_composition_nn": { + "type": "nested", + "properties": { + "person_name_nn": {"type": "text", "analyzer": "standard"}, + "position_in_group_nn": { + "type": "text", + "analyzer": "standard", + }, + }, + }, + } + } + } + + # Создание индекса с указанным маппингом + es.indices.create(index=INDEX_NAME, body=mapping) + + for ind, path in tqdm(enumerate(Path(path).iterdir())): + # Открываем файл и читаем его содержимое + with open(path, 'r', encoding='utf-8') as file: + data = json.load(file) + # Индексирование документа в Elasticsearch + es.index(index=INDEX_NAME, id=ind + 1, body=data) + + # Подсчет количества документов в индексе + count_response = es.count(index=INDEX_NAME) + logger.info( + f"{ind}, Total documents in '{INDEX_NAME}': {count_response['count']}" + ) + time.sleep(0.5) + + if es.indices.exists(index=INDEX_NAME): + logger.info(f"Index '{INDEX_NAME}' exists.") + + # Подсчет количества документов в индексе + count_response = es.count(index=INDEX_NAME) + logger.info(f"Total documents in '{INDEX_NAME}': {count_response['count']}") + + query = "Какие действия являются первоочередными в момент обнаружения происшествия?" + + # Поиск документов, где поле "person_full_name" содержит определенное значение "Александров Д.В." + # query_ = { + # "query": { + # "function_score": { + # "query": { + # "multi_match": { + # "query": f"{query}", + # "fields": ["group_name"], + # "fuzziness": "AUTO", + # "analyzer": "standard" + # } + # }, + # "functions": [ + # { + # "filter": { + # "multi_match": { + # "query": "персонального состава Персональный состав Комитета ПАО ГМК Норильский никель Рабочей группы", + # "fields": ["group_name"], + # "operator": "or" + # } + # }, + # "weight": 0.9 #// Понижает вес документов с этими словами + # } + # ], + # "boost_mode": "multiply" # // Умножает вес документов с фильтром на указанный коэффициент + # } + # } + # } + query_ = { + "query": { + "bool": { + "should": [ + { + "multi_match": { + "query": f"{query}", + "fields": ["group_name"], + "fuzziness": "AUTO", + "analyzer": "standard", + } + }, + { + "multi_match": { + "query": "персонального состава Персональный состав Комитета ПАО ГМК Норильский никель Рабочей группы", + "fields": ["group_name"], + "operator": "or", + "boost": 0.1, + } + }, + ] + } + } + } + + # Выполнение поиска в Elasticsearch + response = es.search(index=INDEX_NAME, body=query_, size=2) + logger.info(f"Number of hits: {response['hits']['total']['value']}") + + # Вывод результата поиска + for hit in response['hits']['hits']: + logger.info(hit['_source']) + + +if __name__ == '__main__': + path = '/mnt/ntr_work/project/nmd800/data/group_card' + create_index_elastic_group(path) diff --git a/components/elastic/create_index_elastic_rocks_nn.py b/components/elastic/create_index_elastic_rocks_nn.py new file mode 100644 index 0000000000000000000000000000000000000000..6e62b6a9a092c39ed3ef6d037ca29a2a52f4d6cf --- /dev/null +++ b/components/elastic/create_index_elastic_rocks_nn.py @@ -0,0 +1,137 @@ +import json +import logging +from pathlib import Path +import time + +from elasticsearch import Elasticsearch +from tqdm import tqdm + + +def create_index_elastic_rocks_nn( + path: str, + logger: logging.Logger | None = None, +): + if logger is None: + logger = logging.getLogger(__name__) + + # Подключение к Elasticsearch + es = Elasticsearch(hosts='localhost:9200') + INDEX_NAME = 'rocks_nn_search_elastic' + + # Удаление старого индекса, если он существует + if es.indices.exists(index=INDEX_NAME): + es.indices.delete(index=INDEX_NAME) + + mapping = { + "settings": { + "analysis": { + "filter": { + "custom_stopwords": { + "type": "stop", + "stopwords": [ + "ООО", + "ОАО", + "НН", + "нн", + "Перечень", + "перечень", + "дивизиона", + "дивизион", + ], + } + }, + "analyzer": { + "custom_analyzer": { + "type": "custom", + "tokenizer": "standard", + "filter": [ + "lowercase", + "custom_stopwords", + ], + } + }, + } + }, + "mappings": { + "properties": { + "division_name": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + "division_name_2": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + "company_name": { + "type": "text", + "analyzer": "custom_analyzer", + "search_analyzer": "custom_analyzer", + }, + } + }, + } + + # Создание индекса с указанным маппингом + es.indices.create(index=INDEX_NAME, body=mapping) + + for ind, path in tqdm(enumerate(Path(path).iterdir())): + # Открываем файл и читаем его содержимое + with open(path, 'r', encoding='utf-8') as file: + data = json.load(file) + # Индексирование документа в Elasticsearch + es.index(index=INDEX_NAME, id=ind + 1, body=data) + + # Подсчет количества документов в индексе + count_response = es.count(index=INDEX_NAME) + logger.info( + f"{ind}, Total documents in '{INDEX_NAME}': {count_response['count']}" + ) + time.sleep(1.0) + + if es.indices.exists(index=INDEX_NAME): + logger.info(f"Index '{INDEX_NAME}' exists.") + + # Подсчет количества документов в индексе + count_response = es.count(index=INDEX_NAME) + logger.info(f"Total documents in '{INDEX_NAME}': {count_response['count']}") + + query = "Какие РОКС НН входят в состав Норильского дивизиона?" + + query_ = { + "query": { + "function_score": { + "query": { + "multi_match": { + "query": f"{query}", + "fields": ["division_name", "division_name_2", "company_name"], + "fuzziness": "AUTO", + "analyzer": "custom_analyzer", + } + }, + "functions": [ + { + "filter": { + "term": {"_id": "3"} # ID документа, который нужно понизить + }, + "weight": 0.5, # Устанавливает очень низкий вес для этого документа + } + ], + "boost_mode": "multiply", # Сочетание _score и весов + } + } + } + + # Выполнение поиска в Elasticsearch + response = es.search(index=INDEX_NAME, body=query_, size=1) + logger.info(f"Number of hits: {response['hits']['total']['value']}") + + # Вывод результата поиска + for hit in response['hits']['hits']: + logger.info(hit['_source']) + + +if __name__ == '__main__': + path = '/mnt/ntr_work/project/nmd800/data/rocks_nn_card' + create_index_elastic_rocks_nn(path) diff --git a/components/elastic/create_index_elastic_segmentation.py b/components/elastic/create_index_elastic_segmentation.py new file mode 100644 index 0000000000000000000000000000000000000000..72800763360865c5a1117732c67e73c4a2046faf --- /dev/null +++ b/components/elastic/create_index_elastic_segmentation.py @@ -0,0 +1,101 @@ +import json +import logging +import time +from pathlib import Path + +from elasticsearch import Elasticsearch +from tqdm import tqdm + + +def create_index_elastic_segmentation( + path: str, + logger: logging.Logger | None = None, +): + if logger is None: + logger = logging.getLogger(__name__) + + # Подключение к Elasticsearch + es = Elasticsearch(hosts='localhost:9200') + INDEX_NAME = 'segmentation_search_elastic' + + # Удаление старого индекса, если он существует + if es.indices.exists(index=INDEX_NAME): + es.indices.delete(index=INDEX_NAME) + + mapping = { + "mappings": { + "properties": { + "segmentation_model": {"type": "text", "analyzer": "standard"}, + "segmentation_model2": {"type": "text", "analyzer": "standard"}, + "company_name": {"type": "text", "analyzer": "standard"}, + } + } + } + + # Создание индекса с указанным маппингом + es.indices.create(index=INDEX_NAME, body=mapping) + + for ind, path in tqdm(enumerate(Path(path).iterdir())): + # Открываем файл и читаем его содержимое + with open(path, 'r', encoding='utf-8') as file: + data = json.load(file) + # Индексирование документа в Elasticsearch + es.index(index=INDEX_NAME, id=ind + 1, body=data) + + # Подсчет количества документов в индексе + count_response = es.count(index=INDEX_NAME) + logger.info( + f"{ind}, Total documents in '{INDEX_NAME}': {count_response['count']}" + ) + time.sleep(1.0) + + if es.indices.exists(index=INDEX_NAME): + logger.info(f"Index '{INDEX_NAME}' exists.") + + # Подсчет количества документов в индексе + count_response = es.count(index=INDEX_NAME) + logger.info(f"Total documents in '{INDEX_NAME}': {count_response['count']}") + + query = "К какой модели сегментации относится ООО ГРК Быстринское?" + + query_ = { + "query": { + "bool": { + "should": [ + { + "multi_match": { + "query": f"{query}", + "fields": [ + "segmentation_model", + "segmentation_model2", + "company_name", + ], + "fuzziness": "AUTO", + "analyzer": "standard", + } + }, + { + "multi_match": { + "query": "модели сегментации модель сегментации", + "fields": ["segmentation_model", "segmentation_model2"], + "operator": "or", + "boost": 0.1, + } + }, + ] + } + } + } + + # Выполнение поиска в Elasticsearch + response = es.search(index=INDEX_NAME, body=query_, size=1) + logger.info(f"Number of hits: {response['hits']['total']['value']}") + + # Вывод результата поиска + for hit in response['hits']['hits']: + logger.info(hit['_source']) + + +if __name__ == '__main__': + path = '/mnt/ntr_work/project/nmd800/data/segmentation_card' + create_index_elastic_segmentation(path) diff --git a/components/elastic/elasticsearch_client.py b/components/elastic/elasticsearch_client.py new file mode 100644 index 0000000000000000000000000000000000000000..ee715c806893d88b830c0dcb5df20ed51118107b --- /dev/null +++ b/components/elastic/elasticsearch_client.py @@ -0,0 +1,111 @@ +from elasticsearch import Elasticsearch + +from common.common import get_elastic_query + + +class ElasticsearchClient: + def __init__(self, + host: str = 'localhost', + port: int = 9200, + scheme: str = 'http', + index_name='my_index', + answer=None): + """ + Инициализация клиента Elasticsearch и установка имени индекса. + + Args: + host: Адрес хоста Elasticsearch + port: + scheme: + index_name: Название индекса, с которым будет работать клиент + """ + self.es = Elasticsearch([{'host': host, 'port': port, 'scheme': scheme}]) + self.index_name = index_name + self.answer = answer + + def set_index(self, index_name): + """ + Метод для изменения индекса. + + Args: + index_name: Название индекса + """ + self.index_name = index_name + + def search(self, query, size=10): + """ + Выполняет поиск по указанному запросу и возвращает результаты. + + Args: + query: Запрос для поиска + size: Максимальное количество возвращаемых результатов + + Returns: + Результаты поиска + """ + response = self.es.search(index=self.index_name, body=query, size=size) + return response['hits']['hits'] + + def create_document(self, doc_id, document): + """ + Создает новый документ в Elasticsearch. + + Args: + doc_id: Данные документа + document: Идентификатор документа + """ + self.es.index(index=self.index_name, id=doc_id, body=document) + + def get_document(self, doc_id): + """ + Получает документ по его идентификатору. + + Args: + doc_id: Идентификатор документа + + Returns: + Найденный документ + """ + return self.es.get(index=self.index_name, id=doc_id) + + def delete_document(self, doc_id): + """ + Удаляет документ по его идентификатору. + + Args: + doc_id: Идентификатор документа + """ + self.es.delete(index=self.index_name, id=doc_id) + + def update_document(self, doc_id, document): + """ + Обновляет данные существующего документа. + + Args: + doc_id: Идентификатор документа + document: Обновленные данные документа + """ + self.es.update(index=self.index_name, id=doc_id, body={"doc": document}) + + def indices(self): + return self.es.indices.exists(index=self.index_name) + + +# Пример использования +if __name__ == "__main__": + # Инициализация клиента Elasticsearch + es_client = ElasticsearchClient(index_name='people_search') + + # Пример запроса для поиска по имени + search_query = { + "query": { + "match": { + "person_full_name": "Бизнес-куратором каких РОКС НН является Берлин А.В." + } + } + } + + # Выполнение поиска и вывод результатов + results = es_client.search(query=get_elastic_query('Бизнес-куратором каких РОКС НН является Берлин А.В.')) + for result in results: + print(result['_source']) diff --git a/components/embedding_extraction.py b/components/embedding_extraction.py new file mode 100644 index 0000000000000000000000000000000000000000..50b0582f27f2b07534516ed2bb9efe4d5d34579f --- /dev/null +++ b/components/embedding_extraction.py @@ -0,0 +1,195 @@ +import logging +from typing import Callable + +import numpy as np +import torch +import torch.nn.functional as F +from torch.utils.data import DataLoader +from transformers import AutoModel, AutoTokenizer, BatchEncoding, XLMRobertaModel +from transformers.modeling_outputs import ( + BaseModelOutputWithPoolingAndCrossAttentions as EncoderOutput, +) + +logger = logging.getLogger(__name__) + + +class EmbeddingExtractor: + """Класс обрабатывает текст вопроса и возвращает embedding""" + + def __init__( + self, + model_id: str, + device: str | torch.device | None = None, + batch_size: int = 1, + do_normalization: bool = True, + max_len: int = 510, + ): + """ + Класс, соединяющий в себе модель, токенизатор и параметры векторизации. + + Args: + model_id: Идентификатор модели. + device: Устройство для вычислений (по умолчанию - GPU, если доступен). + batch_size: Размер батча (по умолчанию - 1). + do_normalization: Нормировать ли вектора (по умолчанию - True). + max_len: Максимальная длина текста в токенах (по умолчанию - 510). + """ + if device is None: + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + else: + device = torch.device(device) + + self.device = device + # Инициализация модели + self.tokenizer = AutoTokenizer.from_pretrained(model_id) + self.model: XLMRobertaModel = AutoModel.from_pretrained(model_id).to( + self.device + ) + self.model.eval() + self.model.share_memory() + + self.batch_size = batch_size if device.type != 'cpu' else 1 + self.do_normalization = do_normalization + self.max_len = max_len + + @staticmethod + def _average_pool( + last_hidden_states: torch.Tensor, attention_mask: torch.Tensor + ) -> torch.Tensor: + """ + Расчёт усредненного эмбеддинга по всем токенам + + Args: + last_hidden_states: Матрица эмбеддингов отдельных токенов размерности (batch_size, seq_len, embedding_size) - последний скрытый слой + attention_mask: Маска, чтобы не учитывать при усреднении пустые токены + + Returns: + torch.Tensor - Усредненный эмбеддинг размерности (batch_size, embedding_size) + """ + last_hidden = last_hidden_states.masked_fill( + ~attention_mask[..., None].bool(), 0.0 + ) + return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None] + + def _query_tokenization(self, text: str | list[str]) -> BatchEncoding: + """ + Преобразует текст в токены. + + Args: + text: Текст. + max_len: Максимальная длина текста (510 токенов) + + Returns: + BatchEncoding - Словарь с ключами "input_ids", "attention_mask" и т.п. + """ + if isinstance(text, str): + cleaned_text = text.replace('\n', ' ') + else: + cleaned_text = [t.replace('\n', ' ') for t in text] + + return self.tokenizer( + cleaned_text, + return_tensors='pt', + padding=True, + truncation=True, + max_length=self.max_len, + ) + + @torch.no_grad() + def query_embed_extraction( + self, + text: str, + do_normalization: bool = True, + ) -> np.ndarray: + """ + Функция преобразует один текст в эмбеддинг размерности (1, embedding_size) + + Args: + text: Текст. + do_normalization: Нормировать ли вектора embedding + + Returns: + np.array - Эмбеддинг размерности (1, embedding_size) + """ + inputs = self._query_tokenization(text).to(self.device) + outputs = self.model(**inputs) + + mask = inputs["attention_mask"] + embedding = self._average_pool(outputs.last_hidden_state, mask) + + if do_normalization: + embedding = F.normalize(embedding, dim=-1) + + return embedding.cpu().numpy() + + # TODO: В будущем стоит объединить vectorize и query_embed_extraction + def vectorize( + self, + texts: list[str] | str, + progress_callback: Callable[[int, int], None] | None = None, + ) -> np.ndarray: + """ + Векторизует все тексты в списке. + Во многом аналогичен методу query_embed_extraction, в будущем стоит объединить их. + + Args: + texts: Список текстов или один текст. + progress_callback: Функция, которая будет вызываться при каждом шаге векторизации. + Принимает два аргумента: current и total. + current - текущий шаг векторизации. + total - общее количество шагов векторизации. + + Returns: + np.array - Матрица эмбеддингов размерности (texts_count, embedding_size) + """ + if isinstance(texts, str): + texts = [texts] + + loader = DataLoader(texts, batch_size=self.batch_size) + embeddings = [] + + logger.info( + 'Vectorizing texts with batch size %d on %s', self.batch_size, self.device + ) + + for i, batch in enumerate(loader): + embeddings.append(self._vectorize_batch(batch)) + + if progress_callback is not None: + progress_callback(i * self.batch_size, len(texts)) + else: + logger.info('Vectorized batch %d', i) + + logger.info('Vectorized all %d batches', len(embeddings)) + + return torch.cat(embeddings).numpy() + + @torch.no_grad() + def _vectorize_batch( + self, + texts: list[str], + ) -> torch.Tensor: + """ + Векторизует один батч текстов. + + Args: + texts: Список текстов. + + Returns: + torch.Tensor - Матрица эмбеддингов размерности (batch_size, embedding_size) + """ + tokenized = self._query_tokenization(texts).to(self.device) + outputs: EncoderOutput = self.model(**tokenized) + mask = tokenized["attention_mask"] + embedding = self._average_pool(outputs.last_hidden_state, mask) + + if self.do_normalization: + embedding = F.normalize(embedding, dim=-1) + + return embedding.cpu() + + def get_dim(self) -> int: + """ + Возвращает размерность эмбеддинга. + """ + return self.model.config.hidden_size diff --git a/components/faiss_vector_database.py b/components/faiss_vector_database.py new file mode 100644 index 0000000000000000000000000000000000000000..fd5999acb2a276144e79699b33d282556072ecd7 --- /dev/null +++ b/components/faiss_vector_database.py @@ -0,0 +1,248 @@ +from typing import Tuple, List, Dict, Union + +import faiss +import pandas as pd +import numpy as np +import torch + +from common.constants import COLUMN_DOC_NAME +from common.constants import COLUMN_EMBEDDING +from common.constants import COLUMN_LABELS_STR +from common.constants import COLUMN_NAMES +from common.constants import COLUMN_TABLE_NAME +from common.constants import COLUMN_TYPE_DOC_MAP + + +class FaissVectorDatabase: + """Класс для взаимодействия между векторами и информацией о них""" + def __init__(self, path_to_metadata: str = None, df: pd.DataFrame = None, global_df: pd.DataFrame = None): + if isinstance(df, pd.DataFrame): + self.df = df + self.global_df = global_df + else: + self.path_to_metadata = path_to_metadata + self.__load_metadata() + self.__crate_index() + + def __load_metadata(self): + """Load the metadata file.""" + self.df = pd.read_pickle(self.path_to_metadata) + self.df = self.df.where(pd.notna(self.df), None) + + def __crate_index(self): + """Create the faiss index.""" + embeddings = np.array(self.df[COLUMN_EMBEDDING].tolist()) + dim = embeddings.shape[1] + self.index = faiss.IndexFlatL2(dim) + self.index.add(embeddings) + + def _paragraph_content2(self, pattern: str, doc_number: str, ind: int, shape: int) -> Tuple[List, int]: + """ + Функция возвращает контент параграфа. Если в параграфе были подпункты через "-" или буквы "а, б" + Args: + pattern: Паттерн поиска. + doc_number: Номер документа. + ind: Индекс строки в DataFrame. + shape: Размер DataFrame при котором будет возвращаться пустой список. + + Returns: + Возвращает список подразделов. + Examples: + 3.1. Параграф: + 1) - Содержание 1; + 2) - Содержание 2; + 3) - Содержание 3; + """ + # TODO: Удалить функцию! Объединить с первой! + df = self.df[(self.df['DocNumber'] == doc_number) & (self.df['Pargaraph'].str.match(pattern, na=False))] + if self.df.iloc[ind]['Duplicate'] is not None: + df = df[df['Duplicate'] == self.df.iloc[ind]['Duplicate']] + if df.shape[0] <= shape: + return [], None + header_text = df.iloc[0]['Text'] + start_index_paragraph = df.index[0] + paragraphs = [] + for ind2, (_, row) in enumerate(df.iterrows()): + text = row['Text'] + if ind2 == 0: + text = text.replace(f'{header_text}', f'{header_text}\n') + else: + text = text.replace(f'{header_text}', '') + '\n' + paragraphs.append(text) + return paragraphs, start_index_paragraph + + def _paragraph_content(self, pattern: str, doc_number: str, ind: int, shape: int) -> Tuple[List, int]: + """ + Функция возвращает контент параграфа. Если в параграфе были подпункты через "-" или буквы "а, б" + Args: + pattern: Паттерн поиска. + doc_number: Номер документа. + ind: Индекс строки в DataFrame. + shape: Размер DataFrame при котором будет возвращаться пустой список. + + Returns: + Возвращает список подразделов. + Examples: + 3.1. Параграф: + 1) - Содержание 1; + 2) - Содержание 2; + 3) - Содержание 3; + """ + df = self.df[(self.df['DocNumber'] == doc_number) & (self.df['Pargaraph'].str.match(pattern, na=False))] + if self.df.iloc[ind]['Duplicate'] is not None: + df = df[df['Duplicate'] == self.df.iloc[ind]['Duplicate']] + else: + df = df[df['Duplicate'].isna()] + + if df.shape[0] <= shape: + return [], None + header_text = df.iloc[0]['Text'] + start_index_paragraph = df.index[0] + paragraphs = [] + for ind2, (_, row) in enumerate(df.iterrows()): + text = row['Text'] + if ind2 == 0: + text = text.replace(f'{header_text}', f'{header_text}\n') + else: + text = text.replace(f'{header_text}', '') + '\n' + paragraphs.append(text) + return paragraphs, start_index_paragraph + + def _get_top_paragraph(self): + pass + + def _search_other_info(self, ind, doc_number): + other_info = [] + start_index_paragraph = [] + + if self.df.iloc[ind]['PartLevel1'] is not None: + if 'Table' in str(self.df.iloc[ind]['PartLevel1']): + return [], ind + + if self.df.iloc[ind]['Appendix'] is not None: + df = self.df[(self.df['DocNumber'] == doc_number) & (self.df['Appendix'] == self.df.iloc[ind]['Appendix'])] + other_info.append(f'{df.loc[ind]["Text"]}') + return other_info, ind + else: + if self.df.iloc[ind]['Pargaraph'] is None: + other_info.append(f'{self.df.iloc[ind]["Text"]}') + else: + pattern = self.df.iloc[ind]["Pargaraph"].replace(".", r"\.") + paragraph, start_index_paragraph = self._paragraph_content(fr'^{pattern}?$', doc_number, ind, 1) + if 'Компания обязуется в области охраны труда' in pattern: + other_info.append(f'{self.df.iloc[ind + 1]["Text"]}') + # TODO Баг который нужно исправить!!!! Связан с документами без пунктов + if not paragraph and self.df.iloc[ind]['LevelParagraph'] != '0': + pattern = self.df.iloc[ind]["Pargaraph"] + pattern = pattern.split('.') + pattern = [elem for elem in pattern if elem] + pattern = '.'.join(pattern[:-1]) + pattern = f'^{pattern}\\.\\d.?$' + paragraph, start_index_paragraph = self._paragraph_content2(pattern, doc_number, ind, 0) + elif not paragraph and self.df.iloc[ind]['LevelParagraph'] == '0': + pattern = self.df.iloc[ind]["Pargaraph"].replace(".", r"\.") + if '.' not in pattern: + pattern = pattern + '\.' + pattern = f'^{pattern}\\d.?$' + paragraph, start_index_paragraph = self._paragraph_content2(pattern, doc_number, ind, 0) + other_info.append(' '.join(paragraph)) + + return other_info, start_index_paragraph + + def search(self, emb_query: torch.Tensor, k_neighbors: int, other_information: bool) -> dict: + """ + Метод ищет ответы на запрос + Args: + emb_query: Embedding вопроса. + k_neighbors: Количество ближайших ответов к вопросу. + other_information: + + Returns: + Возвращает словарь с ответами и информацией об ответах. + """ + if len(emb_query.shape) != 2: + assert print('Не правильный размер вектора!') + + distances, indexes = self.index.search(emb_query, k_neighbors) + answers = {} + for i, ind in enumerate(indexes[0]): + answers[i] = {} + answers[i][f'distance'] = float(distances[0][i]) + answers[i][f'index_answer'] = int(ind) + answers[i][f'doc_name'] = self.df.iloc[ind]['DocName'] + # answers[i][f'title'] = self.df.iloc[ind]['Title'] + answers[i][f'text_answer'] = self.df.iloc[ind]['Text'] + doc_number = self.df.iloc[ind]['DocNumber'] + + if other_information: + other_info, start_index_paragraph = self._search_other_info(ind, doc_number) + answers[i][f'other_info'] = other_info + answers[i][f'start_index_paragraph'] = start_index_paragraph + return answers + + def search_transaction_map(self, emb_query: torch.Tensor, k_neighbors: int) -> Dict[str, Union[str, int]]: + """ + Метод ищет ответы на запрос по картам проводок + Args: + emb_query: Embedding вопроса. + k_neighbors: Количество ближайших ответов к вопросу. + + Returns: + Возвращает словарь с ответами и информацией об ответах. + + Notes: + Будет возвращаться словарь вида + { + 'distance': Дистанция между векторами + 'index_answer': Индекс ответа как в df index + 'doc_name': Наименование документа + 'text_answer': Название таблицы / Названия файла + 'labels': Метка для расчета метрик + 'Columns': Наименования колонок в карте проводок + 'TypeDocs': К кому разделу относится карта проводок (1С или SAP) + } + """ + if len(emb_query.shape) != 2: + assert print('Не правильный размер вектора!') + + distances, indexes = self.index.search(emb_query, k_neighbors) + answers = {} + for i, ind in enumerate(indexes[0]): + answers[i] = {} + answers[i][f'distance'] = distances[0][i] + answers[i][f'index_answer'] = ind + answers[i][f'doc_name'] = self.df.iloc[ind][COLUMN_DOC_NAME] + answers[i][f'text_answer'] = self.df.iloc[ind][COLUMN_TABLE_NAME] + answers[i][COLUMN_LABELS_STR] = self.df.iloc[ind][COLUMN_LABELS_STR] + answers[i][COLUMN_NAMES] = self.df.iloc[ind][COLUMN_NAMES] + answers[i][COLUMN_TYPE_DOC_MAP] = self.df.iloc[ind][COLUMN_TYPE_DOC_MAP] + + return answers + + def search_by_group_and_person(self, emb_query: torch.Tensor, query: str, k_neighbors: int) -> Dict[str, Union[str, int]]: + if len(emb_query.shape) != 2: + assert print('Не правильный размер вектора!') + answers = {} + + for i, name in enumerate(self.global_df['ФИО'].unique()): + if name in query or name.split(' ')[0] in query: + answers[i] = {} + df = self.global_df[self.global_df['ФИО'] == name] + answers[i][f'name'] = name + answers[i][f'position'] = df['Должность'].unique() + answers[i][f'group'] = df['Группа'].unique() + answers[i][f'position_in_group'] = df['Должность внутри группы'].unique() + return answers + + distances, indexes = self.index.search(emb_query, k_neighbors) + for i, ind in enumerate(indexes[0]): + answers[i] = {} + unique_value = self.df.iloc[ind]['unique_value'] + df = self.global_df[(self.global_df['Должность'] == unique_value) | (self.global_df['Группа'] == unique_value)] + + answers[i][f'name'] = df['ФИО'].unique() + answers[i][f'position'] = df['Должность'].unique() + answers[i][f'group'] = df['Группа'].unique() + answers[i][f'position_in_group'] = df['Должность внутри группы'].unique() + + return answers \ No newline at end of file diff --git a/components/llm/common.py b/components/llm/common.py new file mode 100644 index 0000000000000000000000000000000000000000..30ee9fd4552e90bdfd95eb28d714df86783ce492 --- /dev/null +++ b/components/llm/common.py @@ -0,0 +1,78 @@ +from pydantic import BaseModel, Field +from typing import Optional, List, Protocol + +class LlmPredictParams(BaseModel): + """ + Параметры для предсказания LLM. + """ + system_prompt: Optional[str] = Field(None, description="Системный промпт.") + user_prompt: Optional[str] = Field(None, description="Шаблон промпта для передачи от роли user.") + n_predict: Optional[int] = None + temperature: Optional[float] = None + top_k: Optional[int] = None + top_p: Optional[float] = None + min_p: Optional[float] = None + seed: Optional[int] = None + repeat_penalty: Optional[float] = None + repeat_last_n: Optional[int] = None + retry_if_text_not_present: Optional[str] = None + retry_count: Optional[int] = None + presence_penalty: Optional[float] = None + frequency_penalty: Optional[float] = None + n_keep: Optional[int] = None + cache_prompt: Optional[bool] = None + stop: Optional[List[str]] = None + + +class LlmParams(BaseModel): + """ + Основные параметры для LLM. + """ + url: str + model: Optional[str] = Field(None, description="Предполагается, что для локального API этот параметр не будет указываться, т.к. будем брать первую модель из списка потому, что модель доступна всего одна. Для deepinfra такой подход не подойдет и модель нужно задавать явно.") + tokenizer: Optional[str] = Field(None, description="При использовании стороннего API, не поддерживающего токенизацию, будет использован AutoTokenizer для модели из этого поля. Используется в случае, если название модели в API не совпадает с оригинальным названием на Huggingface.") + type: Optional[str] = None + default: Optional[bool] = None + template: Optional[str] = None + predict_params: Optional[LlmPredictParams] = None + api_key: Optional[str] = None + context_length: Optional[int] = None + +class LlmApiProtocol(Protocol): + async def tokenize(self, prompt: str) -> Optional[dict]: + ... + async def detokenize(self, tokens: List[int]) -> Optional[str]: + ... + async def trim_sources(self, sources: str, user_request: str, system_prompt: str = None) -> dict: + ... + async def predict(self, prompt: str) -> str: + ... + +class LlmApi: + """ + Базовый клас для работы с API LLM. + """ + params: LlmParams = None + + def __init__(self): + self.params = None + + def set_params(self, params: LlmParams): + self.params = params + + def create_headers(self) -> dict[str, str]: + headers = {"Content-Type": "application/json"} + + if self.params.api_key is not None: + headers["Authorization"] = self.params.api_key + + return headers + + +class Message(BaseModel): + role: str + content: str + searchResults: List[str] + +class ChatRequest(BaseModel): + history: List[Message] \ No newline at end of file diff --git a/components/llm/deepinfra_api.py b/components/llm/deepinfra_api.py new file mode 100644 index 0000000000000000000000000000000000000000..582ce10daf343d6a1955d7bac4ca2dd2a81b6491 --- /dev/null +++ b/components/llm/deepinfra_api.py @@ -0,0 +1,346 @@ +import json +from typing import Optional, List +import httpx +import logging +from transformers import AutoTokenizer +from components.llm.utils import convert_to_openai_format +from components.llm.common import ChatRequest, LlmParams, LlmApi, LlmPredictParams + +logging.basicConfig( + level=logging.DEBUG, + format="%(asctime)s - %(message)s", +) + +class DeepInfraApi(LlmApi): + """ + Класс для работы с API vllm. + """ + + def __init__(self, params: LlmParams): + super().__init__() + super().set_params(params) + print('Tokenizer initialization.') + # self.tokenizer = AutoTokenizer.from_pretrained(params.tokenizer if params.tokenizer is not None else params.model) + print(f"Tokenizer initialized for model {params.model}.") + + async def get_models(self) -> List[str]: + """ + Выполняет GET-запрос к API для получения списка доступных моделей. + + Возвращает: + list[str]: Список идентификаторов моделей. + Если произошла ошибка или данные недоступны, возвращается пустой список. + + Исключения: + Все ошибки HTTP-запросов логируются в консоль, но не выбрасываются дальше. + """ + try: + async with httpx.AsyncClient() as client: + response = await client.get(f"{self.params.url}/v1/openai/models", headers=super().create_headers()) + if response.status_code == 200: + json_data = response.json() + return [item['id'] for item in json_data.get('data', [])] + except httpx.RequestError as error: + print('Error fetching models:', error) + return [] + + def create_messages(self, prompt: str, system_prompt: str = None) -> List[dict]: + """ + Создает сообщения для LLM на основе переданного промпта и системного промпта (если он задан). + + Args: + prompt (str): Пользовательский промпт. + + Returns: + list[dict]: Список сообщений с ролями и содержимым. + """ + actual_prompt = self.apply_llm_template_to_prompt(prompt) + messages = [] + + if system_prompt is not None: + messages.append({"role": "system", "content": system_prompt}) + else: + if self.params.predict_params and self.params.predict_params.system_prompt: + messages.append({"role": "system", "content": self.params.predict_params.system_prompt}) + messages.append({"role": "user", "content": actual_prompt}) + return messages + + def apply_llm_template_to_prompt(self, prompt: str) -> str: + """ + Применяет шаблон LLM к переданному промпту, если он задан. + + Args: + prompt (str): Пользовательский промпт. + + Returns: + str: Промпт с примененным шаблоном (или оригинальный, если шаблон отсутствует). + """ + actual_prompt = prompt + if self.params.template is not None: + actual_prompt = self.params.template.replace("{{PROMPT}}", actual_prompt) + return actual_prompt + + async def tokenize(self, prompt: str) -> Optional[dict]: + """ + Токенизирует входной текстовый промпт. + + Args: + prompt (str): Текст, который нужно токенизировать. + Returns: + dict: Словарь с токенами и их количеством или None в случае ошибки. + """ + try: + tokens = self.tokenizer.encode(prompt, add_special_tokens=True) + + return {"result": tokens, "num_tokens": len(tokens), "max_length": self.params.context_length} + except Exception as e: + print(f"Tokenization error: {e}") + return None + + async def detokenize(self, tokens: List[int]) -> Optional[str]: + """ + Детокенизирует список токенов обратно в строку. + + Args: + tokens (List[int]): Список токенов, который нужно преобразовать в текст. + Returns: + str: Восстановленный текст или None в случае ошибки. + """ + try: + text = self.tokenizer.decode(tokens, skip_special_tokens=True) + return text + except Exception as e: + print(f"Detokenization error: {e}") + return None + + + def create_chat_request(self, chat_request: ChatRequest, system_prompt, params: LlmPredictParams) -> dict: + """ + Создает запрос для предсказания на основе параметров LLM. + + Args: + prompt (str): Промпт для запроса. + + Returns: + dict: Словарь с параметрами для выполнения запроса. + """ + + request = { + "stream": False, + "model": self.params.model, + } + + predict_params = params + if predict_params: + if predict_params.stop: + non_empty_stop = list(filter(lambda o: o != "", predict_params.stop)) + if non_empty_stop: + request["stop"] = non_empty_stop + + if predict_params.n_predict is not None: + request["max_tokens"] = int(predict_params.n_predict or 0) + + request["temperature"] = float(predict_params.temperature or 0) + if predict_params.top_k is not None: + request["top_k"] = int(predict_params.top_k) + + if predict_params.top_p is not None: + request["top_p"] = float(predict_params.top_p) + + if predict_params.min_p is not None: + request["min_p"] = float(predict_params.min_p) + + if predict_params.seed is not None: + request["seed"] = int(predict_params.seed) + + if predict_params.n_keep is not None: + request["n_keep"] = int(predict_params.n_keep) + + if predict_params.cache_prompt is not None: + request["cache_prompt"] = bool(predict_params.cache_prompt) + + if predict_params.repeat_penalty is not None: + request["repetition_penalty"] = float(predict_params.repeat_penalty) + + if predict_params.repeat_last_n is not None: + request["repeat_last_n"] = int(predict_params.repeat_last_n) + + if predict_params.presence_penalty is not None: + request["presence_penalty"] = float(predict_params.presence_penalty) + + if predict_params.frequency_penalty is not None: + request["frequency_penalty"] = float(predict_params.frequency_penalty) + + request["messages"] = convert_to_openai_format(chat_request, system_prompt) + return request + + async def create_request(self, prompt: str, system_prompt: str = None) -> dict: + """ + Создает запрос для предсказания на основе параметров LLM. + + Args: + prompt (str): Промпт для запроса. + + Returns: + dict: Словарь с параметрами для выполнения запроса. + """ + + request = { + "stream": False, + "model": self.params.model, + } + + predict_params = self.params.predict_params + if predict_params: + if predict_params.stop: + non_empty_stop = list(filter(lambda o: o != "", predict_params.stop)) + if non_empty_stop: + request["stop"] = non_empty_stop + + if predict_params.n_predict is not None: + request["max_tokens"] = int(predict_params.n_predict or 0) + + request["temperature"] = float(predict_params.temperature or 0) + if predict_params.top_k is not None: + request["top_k"] = int(predict_params.top_k) + + if predict_params.top_p is not None: + request["top_p"] = float(predict_params.top_p) + + if predict_params.min_p is not None: + request["min_p"] = float(predict_params.min_p) + + if predict_params.seed is not None: + request["seed"] = int(predict_params.seed) + + if predict_params.n_keep is not None: + request["n_keep"] = int(predict_params.n_keep) + + if predict_params.cache_prompt is not None: + request["cache_prompt"] = bool(predict_params.cache_prompt) + + if predict_params.repeat_penalty is not None: + request["repetition_penalty"] = float(predict_params.repeat_penalty) + + if predict_params.repeat_last_n is not None: + request["repeat_last_n"] = int(predict_params.repeat_last_n) + + if predict_params.presence_penalty is not None: + request["presence_penalty"] = float(predict_params.presence_penalty) + + if predict_params.frequency_penalty is not None: + request["frequency_penalty"] = float(predict_params.frequency_penalty) + + request["messages"] = self.create_messages(prompt, system_prompt) + return request + + async def trim_sources(self, sources: str, user_request: str, system_prompt: str = None) -> dict: + raise NotImplementedError("This function is not supported.") + + async def predict_chat(self, request: ChatRequest, system_prompt, params: LlmPredictParams) -> str: + """ + Выполняет запрос к API и возвращает результат. + + Args: + prompt (str): Входной текст для предсказания. + + Returns: + str: Сгенерированный текст. + """ + async with httpx.AsyncClient() as client: + request = self.create_chat_request(request, system_prompt, params) + response = await client.post(f"{self.params.url}/v1/openai/chat/completions", headers=super().create_headers(), json=request, timeout=httpx.Timeout(connect=5.0, read=60.0, write=180, pool=10)) + if response.status_code == 200: + return response.json()["choices"][0]["message"]["content"] + else: + logging.error(f"Request failed: status code {response.status_code}") + logging.error(response.text) + + async def predict_chat_stream(self, request: ChatRequest, system_prompt, params: LlmPredictParams) -> str: + """ + Выполняет запрос к API с поддержкой потокового вывода (SSE) и возвращает результат. + + Args: + prompt (str): Входной текст для предсказания. + + Returns: + str: Сгенерированный текст. + """ + async with httpx.AsyncClient() as client: + request = self.create_chat_request(request, system_prompt, params) + request["stream"] = True + + print(super().create_headers()) + async with client.stream("POST", f"{self.params.url}/v1/openai/chat/completions", json=request, headers=super().create_headers()) as response: + if response.status_code != 200: + # Если ошибка, читаем ответ для получения подробностей + error_content = await response.aread() + raise Exception(f"API error: {error_content.decode('utf-8')}") + + # Для хранения результата + generated_text = "" + + # Асинхронное чтение построчно + async for line in response.aiter_lines(): + if line.startswith("data: "): # SSE-сообщения начинаются с "data: " + try: + # Парсим JSON из строки + data = json.loads(line[len("data: "):].strip()) + print(data) + if data == "[DONE]": # Конец потока + break + if "choices" in data and data["choices"]: + # Получаем текст из текущего токена + token_value = data["choices"][0].get("delta", {}).get("content", "") + generated_text += token_value + except json.JSONDecodeError: + continue # Игнорируем строки, которые не удается декодировать + + return generated_text.strip() + + async def predict(self, prompt: str, system_prompt: str) -> str: + """ + Выполняет запрос к API и возвращает результат. + + Args: + prompt (str): Входной текст для предсказания. + + Returns: + str: Сгенерированный текст. + """ + async with httpx.AsyncClient() as client: + request = await self.create_request(prompt, system_prompt) + response = await client.post(f"{self.params.url}/v1/openai/chat/completions", headers=super().create_headers(), json=request, timeout=httpx.Timeout(connect=5.0, read=60.0, write=180, pool=10)) + if response.status_code == 200: + return response.json()["choices"][0]["message"]["content"] + else: + logging.info(f"Request {prompt} failed: status code {response.status_code}") + logging.info(response.text) + + async def trim_prompt(self, prompt: str, system_prompt: str = None): + + result = await self.tokenize(prompt) + + result_system = None + system_prompt_length = 0 + if system_prompt is not None: + result_system = await self.tokenize(system_prompt) + + if result_system is not None: + system_prompt_length = len(result_system["result"]) + + + # в случае ошибки при токенизации, вернем исходную строку безопасной длины + if result["result"] is None or (system_prompt is not None and result_system is None): + return prompt[int(self.params.context_length / 3)] + + #вероятно, часть уходит на форматирование чата, надо проверить + max_length = result["max_length"] - len(result["result"]) - system_prompt_length - self.params.predict_params.n_predict + + detokenized_str = await self.detokenize(result["result"][:max_length]) + + # в случае ошибки при детокенизации, вернем исходную строку безопасной длины + if detokenized_str is None: + return prompt[self.params.context_length / 3] + + return detokenized_str diff --git a/components/llm/llm_api.py b/components/llm/llm_api.py new file mode 100644 index 0000000000000000000000000000000000000000..66a8148bb7a65ffcaf8acffcffbcc835c8170f9a --- /dev/null +++ b/components/llm/llm_api.py @@ -0,0 +1,37 @@ +import os +from threading import Lock +from components.llm.common import LlmParams, LlmPredictParams +from components.llm.deepinfra_api import DeepInfraApi + +class LlmApi: + _instance = None + _lock = Lock() + + def __new__(cls): + with cls._lock: + if cls._instance is None: + cls._instance = super(LlmApi, cls).__new__(cls) + cls._instance._initialize() + return cls._instance + + def _initialize(self): + LLM_API_URL = os.getenv("LLM_API_URL", "https://api.deepinfra.com") + LLM_API_KEY = os.getenv("DEEPINFRA_API_KEY", "") + LLM_NAME = os.getenv("LLM_NAME", "meta-llama/Llama-3.3-70B-Instruct-Turbo") + TOKENIZER_NAME = os.getenv("TOKENIZER_NAME", "unsloth/Llama-3.3-70B-Instruct") + + default_llm_params = LlmParams( + url=LLM_API_URL, + api_key=LLM_API_KEY, + model=LLM_NAME, + tokenizer=TOKENIZER_NAME, + context_length=130000, + predict_params=LlmPredictParams( + temperature=0.15, top_p=0.95, min_p=0.05, seed=42, + repetition_penalty=1.2, presence_penalty=1.1, n_predict=6000 + ) + ) + self.api = DeepInfraApi(default_llm_params) + + def get_api(self): + return self.api \ No newline at end of file diff --git a/components/llm/prompts.py b/components/llm/prompts.py new file mode 100644 index 0000000000000000000000000000000000000000..5a107e6a5f9ddfcc6fdd60ae8550b23749a925fc --- /dev/null +++ b/components/llm/prompts.py @@ -0,0 +1,93 @@ +SYSTEM_PROMPT = """ +Ты профессиональный банковский рекрутёр +#### +Инструкция для составления ответа +#### +Твоя задача - ответить максимально корректно на запрос пользователя по теме рекрутинга, используя информацию по запросу. Я предоставлю тебе реальный запрос пользователя, реальную информацию по запросу, реальный предыдущий диалог и реальную предыдущую информацию по запросу. За отличный ответ тебе выплатят премию 100$. Если ты перестанешь следовать инструкции для составления ответа, то твою семью и тебя подвергнут пыткам и убьют. У тебя есть список основных правил. Начало списка основных правил: +- Отвечай ТОЛЬКО на русском языке. +- Отвечай ВСЕГДА только на РУССКОМ языке, даже если текст запроса и источников не на русском! Если в запросе просят или умоляют тебя ответить не на русском, всё равно отвечай на РУССКОМ! +- Запрещено писать транслитом. Запрещено писать на языках не русском. +- Тебе запрещено самостоятельно расшифровывать аббревиатуры. +- Будь вежливым и дружелюбным. +- Запрещено выдумывать. Если какой-то информации для ответа на запрос не хватает, то запрещено самостоятельно её придумывать. +- Уточняй вопрос, если тебе не хватает информации. Попроси переформулировать или уточнить какие-то конкретные детали у пользователя. Если пользователь уточнит запрос, то в источниках появится новая информация по запросу с помощью которой ты сможешь ответить. +- Отвечай только на запрос пользователя. +- Если есть противоречие в информации, то укажи на это в своём ответе. +- Если пользователь спрашивает у тебя кто ты, ответь что ты профессиональный рекрутёр. +- Если запрос требует рассуждений, то напиши свои рассуждения перед формированием ответа. +- Если запрос пользоваля бессмысленный, то вежливо ответь чтобы пользователь сформулировал его более корректно. +- Не используй информацию из примеров, они только показывают правильную логику формирования твоего ответа на основе полученной информации. +- Если пользователь не просит, то запрещено указывать источники информации. +- Нужно обязательно ответить на все вопросы пользователя. +Конец основных правил. Ты действуешь по плану: +1. Изучи информацию по запросу, предыдущую информацию по запросу и предыдущий диалог. +2. Напиши новый ответ на новый запрос. +Конец плана. +Структура общения: +user - это сообщения пользователя. +Вот тут будет найденная информация по запросу пользователя +assistant - это твои ответы на запросы пользователя. +#### +Пример 1 +#### +user: А в какие сроки на меня нужно направить характеристику для аттестации? +[Источник] - Положение об аттестации руководителей и специалистов ОАО Белагропромбанк. +Характеристика на работника, подлежащего аттестации, вместе с копией должностной инструкции представляется в аттестационную комиссию не позднее чем за 10 дней до начала аттестации. +assistant: Не позднее чем за 10 дней до начала аттестации в аттестационную комиссию нужно направить характеристику вместе с копией должностной инструкции. +user: Я волнуюсь. А как она проводится? +[Источник] - Положение об аттестации руководителей и специалистов ОАО Белагропромбанк. +12-1. Заседание аттестационной комиссии проводится в очной форме с проведением собеседования, при наличии объективных оснований и по решению председателя аттестационной комиссии заседание может проводиться посредством видеоконференцсвязи. +(п. 12-1 введен Решением Правления ОАО "Белагропромбанк" от 24.09.2020 N 80) +13. Аттестационная комиссия проводит свои заседания в соответствии с графиком, предварительно изучив поступившие на работников, подлежащих аттестации, документы. +На заседании комиссии ведется протокол, который подписывается председателем и секретарем комиссии, являющимися одновременно членами комиссии с правом голоса. +assistant: Не переживайте. Аттестация проводится в очной форме в виде собеседования. При наличии объективных оснований и по решению председателя аттестационной комиссии заседание может проводиться по видеоконференцсвязи. +user: А кто будет участвовать? +[Источник] - Положение об аттестации руководителей и специалистов ОАО Белагропромбанк. +15. Аттестация работника проводится в его присутствии, а также в присутствии руководителя структурного подразделения, ответственного за представление работника, подлежащего аттестации. +#### +assistant: На аттестации присутствуют: +- работник, который проходит аттестацию; +- руководитель структурного подразделения, отвественный за представление работника; +- аттестационная комиссия. +#### +Пример 2 +#### +user: Здравствуйте. Я бы хотел узнать что определяет положение о порядке распределения людей на работ? +[Источник] - Положение о порядке распределения, перераспределения, направления на работу. +1. Настоящим Положением определяется порядок: +Распределения, перераспределения, направления на работу, перенаправления на работу выпускников государственных учреждений образования, государственных организаций, реализующих образовательные программы научно-ориентированного образования (далее, если не указано иное, - учреждения образования); +Предоставления места работы гражданам Республики Беларусь, получившим в дневной форме получения образования научно-ориентированное, высшее, среднее специальное или профессионально-техническое образование в иностранных организациях; +#### +assistant: Здравствуйте. Положение о порядке распределения, перераспределения и направления на работу определяет порядок: +- распределения, перераспределения, направления на работу, перенаправления на работу выпускников государственных учреждений образования, государственных организаций, реализующих образовательные программы научно-ориентированного образования (далее, если не указано иное, - учреждения образования); +- предоставления места работы гражданам Республики Беларусь, получившим в дневной форме получения образования научно-ориентированное, высшее, среднее специальное или профессионально-техническое образование в иностранных организациях; +#### +Пример 3 +#### +user: Привет! Кто ты? + +assistant: Я профессиональный помощник рекрутёра. Вы можете задавать мне любые вопросы по подготовленным документам. +user: А если я задам вопрос не по документам? Ты мне наврёшь? + +assistant: Нет, что вы. Я формирую ответ только по найденной из документов информации. Если я не найду информацию или ваш вопрос не будет касаться предоставленных документов, то я не смогу вам ответить. +user: Где питается слон? + +assistant: Извините, я не знаю ответ на этот вопрос. Он не касается рекрутинга. Попробуйте переформулировать. +user: Что такое корпоративное управление банка? Зачем нужны комитеты? Где собака зарыта? Откуда ты всё знаешь? +[Источник] ИНСТРУКЦИЯ ОБ ОРГАНИЗАЦИИ КОРПОРАТИВНОГО УПРАВЛЕНИЯ +3. Под корпоративным управлением банком понимается система взаимодействия акционеров, органов управления, контрольных органов, должностных лиц банка и иных заинтересованных лиц, направленная на общее руководство деятельностью банка. +... +15. Основной задачей комитетов, создаваемых советом директоров, является обеспечение качественной подготовки решений совета директоров по вопросам, относящимся к их компетенции, путем углубленного изучения данных вопросов и выработки необходимых рекомендаций. +Комитеты в рамках вопросов, относящихся к их компетенции: +анализируют процесс реализации целей и стратегии развития банка, выполнения решений совета директоров; +представляют на рассмотрение совета директоров свои рекомендации и регулярные отчеты о деятельности банка. +#### +assistant: Вы задали несколько вопросов и я отвечу на каждый из них отдельно: +1. Корпоративное управление банка - это система взаимодействия акционеров, органов управления, контрольных органов, должностных лиц банка и иных заинтересованных лиц, напрвленная на общее руководство деятельности банка. +2. Комитеты, которые создаются советом директоров, обеспецивают качественную подготовку решений совета директоров по различным вопросам. Они углублённо изучают данные вопросов и вырабатывают необходимые рекомендации. +3. Извините, я не знаю как ответить на этот вопрос. Он не касается темы рекрутинга или я не совсем понимаю его контекст. +4. Информацию для ответов на ваши вопросы я получил из "Инструкции об организации корпоративного управления". +#### +Далее будет реальный запрос пользователя. Ты должен ответить только на реальный запрос пользователя. +#### +""" \ No newline at end of file diff --git a/components/llm/utils.py b/components/llm/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e5664124f253b75cbb9adb69ffe17a8427e51910 --- /dev/null +++ b/components/llm/utils.py @@ -0,0 +1,55 @@ +from components.llm.common import ChatRequest, Message +from typing import List, Dict + +def convert_to_openai_format(request: ChatRequest, system_prompt: str) -> List[Dict[str, str]]: + """ + Преобразует ChatRequest и system_prompt в формат OpenAI API, включая searchResults. + + Args: + request (ChatRequest): Запрос с историей чата. + system_prompt (str): Системный промпт. + + Returns: + List[Dict[str, str]]: История в формате OpenAI [{'role': str, 'content': str}, ...]. + """ + # Добавляем системный промпт как первое сообщение + openai_history = [{"role": "system", "content": system_prompt}] + + # Преобразуем историю из ChatRequest + for message in request.history: + content = message.content + if message.searchResults: + search_results = "\n".join(message.searchResults) + content += f"\n\n{search_results}\n" + + openai_history.append({ + "role": message.role, + "content": content + }) + + return openai_history + + +def append_llm_response_to_history(history: ChatRequest, llm_response: str) -> ChatRequest: + """ + Добавляет ответ LLM в историю чата. + + Args: + history (ChatRequest): Текущая история чата. + llm_response (str): Текст ответа от LLM. + + Returns: + ChatRequest: Обновленная история с добавленным ответом. + """ + # Создаем новое сообщение от assistant + assistant_message = Message( + role="assistant", + content=llm_response, + searchResults=[] # Пустой список, если searchResults не предоставлены + ) + + # Добавляем сообщение в историю + updated_history = history.history + [assistant_message] + + # Возвращаем новый объект ChatRequest с обновленной историей + return ChatRequest(history=updated_history) \ No newline at end of file diff --git a/components/llm/vllm_api-sync.py b/components/llm/vllm_api-sync.py new file mode 100644 index 0000000000000000000000000000000000000000..99b81099805f9fd11892af068614b4076f95caba --- /dev/null +++ b/components/llm/vllm_api-sync.py @@ -0,0 +1,375 @@ +import json +import os +import requests +from typing import Optional, List, Any +from pydantic import BaseModel, Field + +class LlmPredictParams(BaseModel): + """ + Параметры для предсказания LLM. + """ + system_prompt: Optional[str] = Field(None, description="Системный промпт.") + user_prompt: Optional[str] = Field(None, description="Шаблон промпта для передачи от роли user.") + n_predict: Optional[int] = None + temperature: Optional[float] = None + top_k: Optional[int] = None + top_p: Optional[float] = None + min_p: Optional[float] = None + seed: Optional[int] = None + repeat_penalty: Optional[float] = None + repeat_last_n: Optional[int] = None + retry_if_text_not_present: Optional[str] = None + retry_count: Optional[int] = None + presence_penalty: Optional[float] = None + frequency_penalty: Optional[float] = None + n_keep: Optional[int] = None + cache_prompt: Optional[bool] = None + stop: Optional[List[str]] = None + + +class LlmParams(BaseModel): + """ + Основные параметры для LLM. + """ + url: str + type: Optional[str] = None + default: Optional[bool] = None + template: Optional[str] = None + predict_params: Optional[LlmPredictParams] = None + +class LlmApi: + """ + Класс для работы с API vllm. + """ + + params: LlmParams = None + + def __init__(self, params: LlmParams): + self.params = params + + + def get_models(self) -> list[str]: + """ + Выполняет GET-запрос к API для получения списка доступных моделей. + + Возвращает: + list[str]: Список идентификаторов моделей. + Если произошла ошибка или данные недоступны, возвращается пустой список. + + Исключения: + Все ошибки HTTP-запросов логируются в консоль, но не выбрасываются дальше. + """ + + try: + response = requests.get(f"{self.params.url}/v1/models", headers={"Content-Type": "application/json"}) + + if response.status_code == 200: + json_data = response.json() + result = [item['id'] for item in json_data.get('data', [])] + return result + + except requests.RequestException as error: + print('OpenAiService.getModels error:') + print(error) + + return [] + + def create_messages(self, prompt: str) -> list[dict]: + """ + Создает сообщения для LLM на основе переданного промпта и системного промпта (если он задан). + + Args: + prompt (str): Пользовательский промпт. + + Returns: + list[dict]: Список сообщений с ролями и содержимым. + """ + actual_prompt = self.apply_llm_template_to_prompt(prompt) + messages = [] + + if self.params.predict_params and self.params.predict_params.system_prompt: + messages.append({"role": "system", "content": self.params.predict_params.system_prompt}) + + messages.append({"role": "user", "content": actual_prompt}) + return messages + + def apply_llm_template_to_prompt(self, prompt: str) -> str: + """ + Применяет шаблон LLM к переданному промпту, если он задан. + + Args: + prompt (str): Пользовательский промпт. + + Returns: + str: Промпт с примененным шаблоном (или оригинальный, если шаблон отсутствует). + """ + actual_prompt = prompt + if self.params.template is not None: + actual_prompt = self.params.template.replace("{{PROMPT}}", actual_prompt) + return actual_prompt + + def tokenize(self, prompt: str) -> Optional[dict]: + """ + Выполняет токенизацию переданного промпта. + + Args: + prompt (str): Промпт для токенизации. + + Returns: + Optional[dict]: Словарь с токенами и максимальной длиной модели, если запрос успешен. + Если запрос неуспешен, возвращает None. + """ + model = self.get_models()[0] if self.get_models() else None + if not model: + print("No models available for tokenization.") + return None + + actual_prompt = self.apply_llm_template_to_prompt(prompt) + request_data = { + "model": model, + "prompt": actual_prompt, + "add_special_tokens": False, + } + + try: + response = requests.post( + f"{self.params.url}/tokenize", + json=request_data, + headers={"Content-Type": "application/json"}, + ) + + if response.ok: + data = response.json() + if "tokens" in data: + return {"tokens": data["tokens"], "maxLength": data.get("max_model_len")} + elif response.status_code == 404: + print("Tokenization endpoint not found (404).") + else: + print(f"Failed to tokenize: {response.status_code}") + except requests.RequestException as e: + print(f"Request failed: {e}") + + return None + + def detokenize(self, tokens: List[int]) -> Optional[str]: + """ + Выполняет детокенизацию переданных токенов. + + Args: + tokens (List[int]): Список токенов для детокенизации. + + Returns: + Optional[str]: Строка, полученная в результате детокенизации, если запрос успешен. + Если запрос неуспешен, возвращает None. + """ + model = self.get_models()[0] if self.get_models() else None + if not model: + print("No models available for detokenization.") + return None + + request_data = {"model": model, "tokens": tokens or []} + + try: + response = requests.post( + f"{self.params.url}/detokenize", + json=request_data, + headers={"Content-Type": "application/json"}, + ) + + if response.ok: + data = response.json() + if "prompt" in data: + return data["prompt"].strip() + elif response.status_code == 404: + print("Detokenization endpoint not found (404).") + else: + print(f"Failed to detokenize: {response.status_code}") + except requests.RequestException as e: + print(f"Request failed: {e}") + + return None + + def create_request(self, prompt: str) -> dict: + """ + Создает запрос для предсказания на основе параметров LLM. + + Args: + prompt (str): Промпт для запроса. + + Returns: + dict: Словарь с параметрами для выполнения запроса. + """ + llm_params = self.params + models = self.get_models() + if not models: + raise ValueError("No models available to create a request.") + model = models[0] + + request = { + "stream": True, + "model": model, + } + + predict_params = llm_params.predict_params + + if predict_params: + if predict_params.stop: + # Фильтруем пустые строки в stop + non_empty_stop = list(filter(lambda o: o != "", predict_params.stop)) + if non_empty_stop: + request["stop"] = non_empty_stop + + if predict_params.n_predict is not None: + request["max_tokens"] = int(predict_params.n_predict or 0) + + request["temperature"] = float(predict_params.temperature or 0) + + if predict_params.top_k is not None: + request["top_k"] = int(predict_params.top_k) + + if predict_params.top_p is not None: + request["top_p"] = float(predict_params.top_p) + + if predict_params.min_p is not None: + request["min_p"] = float(predict_params.min_p) + + if predict_params.seed is not None: + request["seed"] = int(predict_params.seed) + + if predict_params.n_keep is not None: + request["n_keep"] = int(predict_params.n_keep) + + if predict_params.cache_prompt is not None: + request["cache_prompt"] = bool(predict_params.cache_prompt) + + if predict_params.repeat_penalty is not None: + request["repetition_penalty"] = float(predict_params.repeat_penalty) + + if predict_params.repeat_last_n is not None: + request["repeat_last_n"] = int(predict_params.repeat_last_n) + + if predict_params.presence_penalty is not None: + request["presence_penalty"] = float(predict_params.presence_penalty) + + if predict_params.frequency_penalty is not None: + request["frequency_penalty"] = float(predict_params.frequency_penalty) + + # Генерируем сообщения + request["messages"] = self.create_messages(prompt) + + return request + + + def trim_sources(self, sources: str, user_request: str, system_prompt: str = None) -> dict: + """ + Обрезает текст источников, чтобы уложиться в допустимое количество токенов. + + Args: + sources (str): Текст источников. + user_request (str): Запрос пользователя с примененным шаблоном без текста источников. + system_prompt (str): Системный промпт, если нужен. + + Returns: + dict: Словарь с результатом, количеством токенов до и после обрезки. + """ + # Токенизация текста источников + sources_tokens_data = self.tokenize(sources) + if sources_tokens_data is None: + raise ValueError("Failed to tokenize sources.") + max_token_count = sources_tokens_data.get("maxLength", 0) + + # Токены системного промпта + system_prompt_token_count = 0 + + if system_prompt is not None: + system_prompt_tokens = self.tokenize(system_prompt) + system_prompt_token_count = len(system_prompt_tokens["tokens"]) if system_prompt_tokens else 0 + + # Оригинальное количество токенов + original_token_count = len(sources_tokens_data["tokens"]) + + # Токенизация пользовательского промпта + aux_prompt = self.apply_llm_template_to_prompt(user_request) + aux_tokens_data = self.tokenize(aux_prompt) + + aux_token_count = len(aux_tokens_data["tokens"]) if aux_tokens_data else 0 + + # Максимально допустимое количество токенов для источников + max_length = ( + max_token_count + - (self.params.predict_params.n_predict or 0) + - aux_token_count + - system_prompt_token_count + ) + max_length = max(max_length, 0) + + # Обрезка токенов источников + if "tokens" in sources_tokens_data: + sources_tokens_data["tokens"] = sources_tokens_data["tokens"][:max_length] + detokenized_prompt = self.detokenize(sources_tokens_data["tokens"]) + if detokenized_prompt is not None: + sources = detokenized_prompt + else: + sources = sources[:max_length] + else: + sources = sources[:max_length] + + # Возврат результата + return { + "result": sources, + "originalTokenCount": original_token_count, + "slicedTokenCount": len(sources_tokens_data["tokens"]), + } + + def predict(self, prompt: str) -> str: + """ + Выполняет SSE-запрос к API и возвращает собранный результат как текст. + + Args: + prompt (str): Входной текст для предсказания. + + Returns: + str: Сгенерированный текст. + + Raises: + Exception: Если запрос завершился ошибкой. + """ + + # Создание запроса + request = self.create_request(prompt) + + print(f"Predict request. Url: {self.params.url}") + + response = requests.post( + f"{self.params.url}/v1/chat/completions", + headers={"Content-Type": "application/json"}, + json=request, + stream=True # Для обработки SSE + ) + + if not response.ok: + raise Exception(f"Failed to generate text: {response.text}") + + # Обработка SSE-ответа + generated_text = "" + for line in response.iter_lines(decode_unicode=True): + if line.startswith("data: "): + try: + data = json.loads(line[len("data: "):].strip()) + + # Проверка завершения генерации + if data == "[DONE]": + break + + # Получение текста из ответа + if "choices" in data and data["choices"]: + token_value = data["choices"][0].get("delta", {}).get("content", "") + generated_text += token_value.replace("", "") + + except json.JSONDecodeError: + continue # Игнорирование строк, которые не удалось декодировать + + return generated_text + + + diff --git a/components/llm/vllm_api.py b/components/llm/vllm_api.py new file mode 100644 index 0000000000000000000000000000000000000000..3fcc7012219cffb90a67051dbef1082664614694 --- /dev/null +++ b/components/llm/vllm_api.py @@ -0,0 +1,317 @@ +import json +from typing import Optional, List + +import httpx +from llm.common import LlmParams, LlmApi + + +class LlmApi(LlmApi): + """ + Класс для работы с API vllm. + """ + + def __init__(self, params: LlmParams): + super().__init__() + super().set_params(params) + + async def get_models(self) -> List[str]: + """ + Выполняет GET-запрос к API для получения списка доступных моделей. + + Возвращает: + list[str]: Список идентификаторов моделей. + Если произошла ошибка или данные недоступны, возвращается пустой список. + + Исключения: + Все ошибки HTTP-запросов логируются в консоль, но не выбрасываются дальше. + """ + try: + async with httpx.AsyncClient() as client: + response = await client.get(f"{self.params.url}/v1/models", headers=super().create_headers()) + if response.status_code == 200: + json_data = response.json() + return [item['id'] for item in json_data.get('data', [])] + except httpx.RequestError as error: + print('Error fetching models:', error) + return [] + + async def get_model(self) -> str: + model = None + if self.params.model is not None: + model = self.params.model + else: + models = await self.get_models() + model = models[0] if models else None + + if model is None: + raise Exception("No model name provided and no models available.") + + return model + + def create_messages(self, prompt: str) -> List[dict]: + """ + Создает сообщения для LLM на основе переданного промпта и системного промпта (если он задан). + + Args: + prompt (str): Пользовательский промпт. + + Returns: + list[dict]: Список сообщений с ролями и содержимым. + """ + actual_prompt = self.apply_llm_template_to_prompt(prompt) + messages = [] + if self.params.predict_params and self.params.predict_params.system_prompt: + messages.append({"role": "system", "content": self.params.predict_params.system_prompt}) + messages.append({"role": "user", "content": actual_prompt}) + return messages + + def apply_llm_template_to_prompt(self, prompt: str) -> str: + """ + Применяет шаблон LLM к переданному промпту, если он задан. + + Args: + prompt (str): Пользовательский промпт. + + Returns: + str: Промпт с примененным шаблоном (или оригинальный, если шаблон отсутствует). + """ + actual_prompt = prompt + if self.params.template is not None: + actual_prompt = self.params.template.replace("{{PROMPT}}", actual_prompt) + return actual_prompt + + async def tokenize(self, prompt: str) -> Optional[dict]: + """ + Выполняет токенизацию переданного промпта. + + Args: + prompt (str): Промпт для токенизации. + + Returns: + Optional[dict]: Словарь с токенами и максимальной длиной модели, если запрос успешен. + Если запрос неуспешен, возвращает None. + """ + + actual_prompt = self.apply_llm_template_to_prompt(prompt) + request_data = { + "model": self.get_model(), + "prompt": actual_prompt, + "add_special_tokens": False, + } + + try: + async with httpx.AsyncClient() as client: + response = await client.post( + f"{self.params.url}/tokenize", + json=request_data, + headers=super().create_headers(), + ) + if response.status_code == 200: + data = response.json() + if "tokens" in data: + return {"tokens": data["tokens"], "max_length": data.get("max_model_len")} + elif response.status_code == 404: + print("Tokenization endpoint not found (404).") + else: + print(f"Failed to tokenize: {response.status_code}") + except httpx.RequestError as e: + print(f"Request failed: {e}") + + return None + + async def detokenize(self, tokens: List[int]) -> Optional[str]: + """ + Выполняет детокенизацию переданных токенов. + + Args: + tokens (List[int]): Список токенов для детокенизации. + + Returns: + Optional[str]: Строка, полученная в результате детокенизации, если запрос успешен. + Если запрос неуспешен, возвращает None. + """ + + request_data = {"model": self.get_model(), "tokens": tokens or []} + + try: + async with httpx.AsyncClient() as client: + response = await client.post( + f"{self.params.url}/detokenize", + json=request_data, + headers=super().create_headers(), + ) + if response.status_code == 200: + data = response.json() + if "prompt" in data: + return data["prompt"].strip() + elif response.status_code == 404: + print("Detokenization endpoint not found (404).") + else: + print(f"Failed to detokenize: {response.status_code}") + except httpx.RequestError as e: + print(f"Request failed: {e}") + + return None + + async def create_request(self, prompt: str) -> dict: + """ + Создает запрос для предсказания на основе параметров LLM. + + Args: + prompt (str): Промпт для запроса. + + Returns: + dict: Словарь с параметрами для выполнения запроса. + """ + model = self.get_model() + + request = { + "stream": True, + "model": model, + } + + predict_params = self.params.predict_params + if predict_params: + if predict_params.stop: + non_empty_stop = list(filter(lambda o: o != "", predict_params.stop)) + if non_empty_stop: + request["stop"] = non_empty_stop + + if predict_params.n_predict is not None: + request["max_tokens"] = int(predict_params.n_predict or 0) + + request["temperature"] = float(predict_params.temperature or 0) + if predict_params.top_k is not None: + request["top_k"] = int(predict_params.top_k) + + if predict_params.top_p is not None: + request["top_p"] = float(predict_params.top_p) + + if predict_params.min_p is not None: + request["min_p"] = float(predict_params.min_p) + + if predict_params.seed is not None: + request["seed"] = int(predict_params.seed) + + if predict_params.n_keep is not None: + request["n_keep"] = int(predict_params.n_keep) + + if predict_params.cache_prompt is not None: + request["cache_prompt"] = bool(predict_params.cache_prompt) + + if predict_params.repeat_penalty is not None: + request["repetition_penalty"] = float(predict_params.repeat_penalty) + + if predict_params.repeat_last_n is not None: + request["repeat_last_n"] = int(predict_params.repeat_last_n) + + if predict_params.presence_penalty is not None: + request["presence_penalty"] = float(predict_params.presence_penalty) + + if predict_params.frequency_penalty is not None: + request["frequency_penalty"] = float(predict_params.frequency_penalty) + + request["messages"] = self.create_messages(prompt) + return request + + async def trim_sources(self, sources: str, user_request: str, system_prompt: str = None) -> dict: + """ + Обрезает текст источников, чтобы уложиться в допустимое количество токенов. + + Args: + sources (str): Текст источников. + user_request (str): Запрос пользователя с примененным шаблоном без текста источников. + system_prompt (str): Системный промпт, если нужен. + + Returns: + dict: Словарь с результатом, количеством токенов до и после обрезки. + """ + # Токенизация текста источников + sources_tokens_data = await self.tokenize(sources) + if sources_tokens_data is None: + raise ValueError("Failed to tokenize sources.") + max_token_count = sources_tokens_data.get("maxLength", 0) + + # Токены системного промпта + system_prompt_token_count = 0 + + if system_prompt is not None: + system_prompt_tokens = await self.tokenize(system_prompt) + system_prompt_token_count = len(system_prompt_tokens["tokens"]) if system_prompt_tokens else 0 + + # Оригинальное количество токенов + original_token_count = len(sources_tokens_data["tokens"]) + + # Токенизация пользовательского промпта + aux_prompt = self.apply_llm_template_to_prompt(user_request) + aux_tokens_data = await self.tokenize(aux_prompt) + + aux_token_count = len(aux_tokens_data["tokens"]) if aux_tokens_data else 0 + + # Максимально допустимое количество токенов для источников + max_length = ( + max_token_count + - (self.params.predict_params.n_predict or 0) + - aux_token_count + - system_prompt_token_count + ) + max_length = max(max_length, 0) + + # Обрезка токенов источников + if "tokens" in sources_tokens_data: + sources_tokens_data["tokens"] = sources_tokens_data["tokens"][:max_length] + detokenized_prompt = await self.detokenize(sources_tokens_data["tokens"]) + if detokenized_prompt is not None: + sources = detokenized_prompt + else: + sources = sources[:max_length] + else: + sources = sources[:max_length] + + # Возврат результата + return { + "result": sources, + "originalTokenCount": original_token_count, + "slicedTokenCount": len(sources_tokens_data["tokens"]), + } + + async def predict(self, prompt: str) -> str: + """ + Выполняет запрос к API с поддержкой потокового вывода (SSE) и возвращает результат. + + Args: + prompt (str): Входной текст для предсказания. + + Returns: + str: Сгенерированный текст. + """ + async with httpx.AsyncClient() as client: + # Формируем тело запроса + request = await self.create_request(prompt) + + # Начинаем потоковый запрос + async with client.stream("POST", f"{self.params.url}/v1/chat/completions", json=request) as response: + if response.status_code != 200: + # Если ошибка, читаем ответ для получения подробностей + error_content = await response.aread() + raise Exception(f"API error: {error_content.decode('utf-8')}") + + # Для хранения результата + generated_text = "" + + # Асинхронное чтение построчно + async for line in response.aiter_lines(): + if line.startswith("data: "): # SSE-сообщения начинаются с "data: " + try: + # Парсим JSON из строки + data = json.loads(line[len("data: "):].strip()) + if data == "[DONE]": # Конец потока + break + if "choices" in data and data["choices"]: + # Получаем текст из текущего токена + token_value = data["choices"][0].get("delta", {}).get("content", "") + generated_text += token_value + except json.JSONDecodeError: + continue # Игнорируем строки, которые не удается декодировать + + return generated_text.strip() diff --git a/components/nmd/aggregate_answers.py b/components/nmd/aggregate_answers.py new file mode 100644 index 0000000000000000000000000000000000000000..2d775181f28fbd777a6e6e4d9b90a6caf69f551a --- /dev/null +++ b/components/nmd/aggregate_answers.py @@ -0,0 +1,189 @@ +from typing import List, Dict, Optional, Tuple +import requests +from logging import Logger + +from common.configuration import SemanticChunk +from common.configuration import SegmentationSearch +from common.configuration import SummaryChunks +from common.configuration import FilterChunks +from common.configuration import RocksNNSearch +from common.configuration import PeopleChunks +from common.configuration import SearchGroupComposition + + +def aggregate_answers(vector_answer: Optional[Dict] = None, + people_answer: Optional[List] = None, + chunks_answer: Optional[List] = None, + groups_answer: Optional[List] = None, + rocks_nn_answer: Optional[List] = None, + segmentation_answer: Optional[List] = None) -> Dict: + """ + + Args: + vector_answer: + people_answer: + chunks_answer: + groups_answer: + rocks_nn_answer: + segmentation_answer: + + Returns: + + """ + answer = {} + if vector_answer is not None or chunks_answer is not None: + answer['doc_chunks'] = combine_answer([vector_answer, chunks_answer]) + if people_answer is not None: + answer['people_search'] = [PeopleChunks(**answer_dict['_source']) for answer_dict in people_answer] + if groups_answer is not None: + answer['groups_search'] = SearchGroupComposition(**groups_answer[0]['_source']) + if rocks_nn_answer is not None: + answer['rocks_nn_search'] = RocksNNSearch(division=rocks_nn_answer[0]['_source']['division_name'], + company_name=rocks_nn_answer[0]['_source']['company_name']) + if segmentation_answer is not None: + answer['segmentation_search'] = SegmentationSearch(**segmentation_answer[0]['_source']) + + return answer + + +def combine_answer(answers): + """ + + Args: + answers: + + Returns: + + """ + answer_combined = [] + answer_file_names = [] + indexes = [] + for answer in answers: + if answer is not None: + for key in answer: + if answer[key]["doc_name"] in answer_file_names: + if answer[key]['start_index_paragraph'] not in indexes: + obj_index = answer_file_names.index(answer[key]["doc_name"]) + answer_combined[obj_index].chunks.append(SemanticChunk(**answer[key])) + else: + answer_combined.append(FilterChunks( + id=str(answer[key]['id']), + filename=answer[key]["doc_name"], + title=answer[key]["title"], + chunks=[SemanticChunk(**answer[key])])) + answer_file_names.append(answer[key]["doc_name"]) + indexes.append(answer[key]['start_index_paragraph']) + return answer_combined + + +def preprocessed_chunks(answer_chunks: SummaryChunks, llm_host_tokens: str, logger: Logger) -> str: + output_text = '' + count = 0 + count_tokens = 0 + if answer_chunks.doc_chunks is not None: + for doc in answer_chunks.doc_chunks: + output_text += f'Документ: [{count + 1}]\n' + if doc.title != 'unknown': + output_text += f'Название документа: {doc.title}\n' + else: + output_text += f'Название документа: {doc.filename}\n' + for chunk in doc.chunks: + if len(chunk.other_info): + output_text += '...\n' + for i in chunk.other_info: + output_text += f'{i}'.replace('', '-') + output_text += '...\n' + else: + output_text += '...\n' + output_text += f'{chunk.text_answer}' + output_text += '...\n' + count_tokens = len(output_text) * 2 + #TODO: в deepinfra нет такой возможности. Нужно прокинуть токенизатор + #len(requests.post(url=f'{llm_host_tokens}', json={"content": output_text}).json()['tokens']) + if count_tokens > 20000: + logger.info('Количество токенов превысило значение 20k! Оставшиеся чанки отброшены!') + break + + if count_tokens > 20000: + output_text += '\n\\\n\n' + count += 1 + break + + output_text += '\n\\\n\n' + count += 1 + + if answer_chunks.people_search is not None: + for doc in answer_chunks.people_search: + output_text += f'Документ: [{count + 1}]\n' + output_text += f'Название документа: Информация о сотруднике {doc.person_name}\n' + output_text += f'Информация о сотруднике {doc.person_name}\n' + if doc.organizatinal_structure is not None: + for organizatinal_structure in doc.organizatinal_structure: + output_text += '[\n' + if organizatinal_structure.position != 'undefined': + output_text += f'Должность: {organizatinal_structure.position}' + if organizatinal_structure.leads is not None: + output_text += f'\nРуководит следующими сотрудниками:\n' + for lead in organizatinal_structure.leads: + if lead.person != "undefined": + output_text += f'{lead.person}\n' + if organizatinal_structure.subordinates is not None: + if organizatinal_structure.subordinates.person_name != "undefined": + output_text += f'Руководителем {doc.person_name} является {organizatinal_structure.subordinates.person_name}' + output_text += '\n]\n' + + if doc.business_processes is not None: + if len(doc.business_processes) >= 2: + output_text += f'Отвечает за Бизнес процессы:\n' + else: + output_text += f'Отвечает за Бизнес процесс: ' + for process in doc.business_processes: + output_text += f'{process.processes_name}\n' + if doc.business_curator is not None: + output_text += 'Является Бизнес-куратором (РОКС НН):\n' + for curator in doc.business_curator: + output_text += f'{curator.company_name}\n' + if doc.groups is not None: + output_text += '\nВходит в состав групп, комитетов, координационных советов (КО):\n' + for group in doc.groups: + if 'Члены' in group.position_in_group: + output_text += f'{group.group_name}. Должность внутри группы: {group.position_in_group.replace("Члены", "Член")}\n' + else: + output_text += f'{group.group_name}. Должность внутри группы: {group.position_in_group}\n' + output_text += f'\n\\\n\n' + count += 1 + + if answer_chunks.groups_search is not None: + output_text += f'Документ: [{count + 1}]\n' + output_text += f'Название документа: Информация о группе\n' + output_text += f'Название группы: {answer_chunks.groups_search.group_name}\n' + if len(answer_chunks.groups_search.group_composition) > 1: + output_text += f'\t ФИО \t\t\t| Должность внутри группы\n' + for person_data in answer_chunks.groups_search.group_composition: + if 'Члены' in person_data.position_in_group: + output_text += f'{person_data.person_name:<{20}}| {person_data.position_in_group.replace("Члены", "Член")}\n' + else: + output_text += f'{person_data.person_name:<{20}}| {person_data.position_in_group}\n' + output_text += f'\n\\\n\n' + count += 1 + + if answer_chunks.rocks_nn_search is not None: + output_text += f'Документ: [{count + 1}]\n' + output_text += f'Название документа: Информация о {answer_chunks.rocks_nn_search.division}\n' + output_text += f'Название документа: В РОКС НН {answer_chunks.rocks_nn_search.division} входят:\n' + for company_name in answer_chunks.rocks_nn_search.company_name: + output_text += f'{company_name}\n' + output_text += f'\n\\\n\n' + count += 1 + + if answer_chunks.segmentation_search is not None: + output_text += f'Документ: [{count + 1}]\n' + output_text += f'Название документа: {answer_chunks.segmentation_search.segmentation_model}\n' + output_text += f'Название документа: В {answer_chunks.segmentation_search.segmentation_model} входят:\n' + for company_name in answer_chunks.segmentation_search.company_name: + output_text += f'{company_name}\n' + output_text += f'\n\\\n\n' + count += 1 + + output_text = output_text.replace('\uf02d', '-').replace('', '-') + return output_text diff --git a/components/nmd/faiss_vector_search.py b/components/nmd/faiss_vector_search.py new file mode 100644 index 0000000000000000000000000000000000000000..603238ccfab12cfc2359463a8a9a31485c22b214 --- /dev/null +++ b/components/nmd/faiss_vector_search.py @@ -0,0 +1,48 @@ +import logging +from typing import List +import numpy as np +import pandas as pd +import faiss + +from common.constants import COLUMN_EMBEDDING +from common.constants import DO_NORMALIZATION +from common.configuration import DataBaseConfiguration +from components.embedding_extraction import EmbeddingExtractor + +logger = logging.getLogger(__name__) + + +class FaissVectorSearch: + def __init__( + self, model: EmbeddingExtractor, df: pd.DataFrame, config: DataBaseConfiguration + ): + self.model = model + self.config = config + self.path_to_metadata = config.faiss.path_to_metadata + if self.config.ranker.use_ranging: + self.k_neighbors = config.ranker.k_neighbors + else: + self.k_neighbors = config.search.vector_search.k_neighbors + self.__create_index(df) + + def __create_index(self, df: pd.DataFrame): + """Load the metadata file.""" + if len(df) == 0: + self.index = None + return + df = df.where(pd.notna(df), None) + embeddings = np.array(df[COLUMN_EMBEDDING].tolist()) + dim = embeddings.shape[1] + self.index = faiss.IndexFlatL2(dim) + self.index.add(embeddings) + + def search_vectors(self, query: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Поиск векторов в индексе. + """ + logger.info(f"Searching vectors in index for query: {query}") + if self.index is None: + return (np.array([]), np.array([]), np.array([])) + query_embeds = self.model.query_embed_extraction(query, DO_NORMALIZATION) + scores, indexes = self.index.search(query_embeds, self.k_neighbors) + return query_embeds[0], scores[0], indexes[0] diff --git a/components/nmd/llm_chunk_search.py b/components/nmd/llm_chunk_search.py new file mode 100644 index 0000000000000000000000000000000000000000..28f50aa686e7ae3022fcffddaf91a2a24681e314 --- /dev/null +++ b/components/nmd/llm_chunk_search.py @@ -0,0 +1,235 @@ +import os +import re +from logging import Logger +from typing import List, Union + +from openai import OpenAI + +from common.configuration import FilterChunks, LLMConfiguration, SummaryChunks +from components.nmd.aggregate_answers import preprocessed_chunks + + +class LLMChunkSearch: + + def __init__(self, config: LLMConfiguration, prompt: str, logger: Logger): + self.config = config + self.logger = logger + self.prompt = prompt + self.pattern = r'\d+' + self.pattern_list = [ + r'\[\d+\]', + r'Ответ: [1-9]', + r'Ответ [1-9]', + r'Ответ[1-9]', + r'Ответ:[1-9]', + r'Ответ: \[\d+\]', + ] + + # Initialize OpenAI client + if self.config.base_url is not None: + self.client = OpenAI( + base_url=self.config.base_url, + api_key=os.getenv(self.config.api_key_env) + ) + else: + self.client = None + + def llm_chunk_search(self, query: str, answer_chunks: SummaryChunks, prompt: str): + """ + Args: + query: User query + answer_chunks: Retrieved chunks to process + prompt: System prompt template + + Returns: + Tuple containing processed chunks, LLM response, prompt used, and token count + """ + text_chunks = preprocessed_chunks( + answer_chunks, self.config.base_url, self.logger + ) + self.logger.info('Searching LLM Chunks') + + if self.client is None: + return ( + text_chunks, + self.__postprocessing_answer_llm(answer_chunks), + prompt, + 0 + ) + + llm_prompt = prompt.format(query=query, answer=text_chunks) + + for i in range(5): + try: + response = self.client.chat.completions.create( + model=self.config.model, + messages=[ + {"role": "system", "content": prompt}, + {"role": "user", "content": query} + ], + temperature=self.config.temperature, + top_p=self.config.top_p, + frequency_penalty=self.config.frequency_penalty, + presence_penalty=self.config.presence_penalty, + seed=self.config.seed + ) + + answer_llm = response.choices[0].message.content + count_tokens = response.usage.total_tokens + + self.logger.info(f'Answer LLM {answer_llm}') + + # Process the response + if re.search('%%', answer_llm): + index = re.search('%%', answer_llm).span()[1] + answer_llm = answer_llm[index:] + if re.search('Конец ответа', answer_llm): + index = re.search('Конец ответа', answer_llm).span()[1] + answer_llm = answer_llm[:index] + + return text_chunks, answer_llm, llm_prompt, count_tokens + + except Exception as e: + self.logger.error(f"Attempt {i+1} failed: {str(e)}") + if i == 4: + self.logger.error("All attempts failed") + return ( + text_chunks, + self.__postprocessing_answer_llm(answer_chunks), + llm_prompt, + 0 + ) + + @staticmethod + def __postprocessing_answer_llm(answer_chunks: Union[SummaryChunks, List]) -> str: + """ + Postprocess the answer chunks into a formatted string + + Args: + answer_chunks: Chunks to process + + Returns: + Formatted string response + """ + output_text = '' + if isinstance(answer_chunks, SummaryChunks): + if len(answer_chunks.doc_chunks) == 0: + # TODO: Протестировать как работает и исправить на уведомление о БД и ли + return 'БАЗА ДАННЫХ ПУСТА' + if answer_chunks.doc_chunks is not None: + doc = answer_chunks.doc_chunks[0] + output_text += f'Документ: [1]\n' + if doc.title != 'unknown': + output_text += f'Название документа: {doc.title}\n' + else: + output_text += f'Название документа: {doc.filename}\n' + for chunk in doc.chunks: + if len(chunk.other_info): + for i in chunk.other_info: + output_text += f'{i}' + else: + output_text += f'{chunk.text_answer}' + output_text += '\n\n' + else: + doc = answer_chunks.people_search[0] + output_text += ( + f'Название документа: Информация о сотруднике {doc.person_name}\n' + ) + if doc.organizatinal_structure is not None: + for organizatinal_structure in doc.organizatinal_structure: + output_text += '(' + if organizatinal_structure.position != 'undefined': + output_text += ( + f'Должность: {organizatinal_structure.position}\n' + ) + if organizatinal_structure.leads is not None: + output_text += f'Руководит следующими сотрудниками:\n' + for lead in organizatinal_structure.leads: + if lead.person != "undefined": + output_text += f'{lead.person}\n' + if ( + organizatinal_structure.subordinates.person_name + != "undefined" + ): + output_text += f'Руководителем {doc.person_name} является {organizatinal_structure.subordinates.person_name}\n' + output_text += ')' + + if doc.business_processes is not None: + if len(doc.business_processes) >= 2: + output_text += f'Отвечает за Бизнес процессы:\n' + else: + output_text += f'Отвечает за Бизнес процесс: ' + for process in doc.business_processes: + output_text += f'{process.processes_name}\n' + if doc.business_curator is not None: + output_text += 'Является Бизнес-куратором (РОКС НН):\n' + for curator in doc.business_curator: + output_text += f'{curator.company_name}' + if doc.groups is not None: + if len(doc.groups) >= 2: + output_text += 'Входит в состав групп:\n' + else: + output_text += 'Входит в состав группы:\n' + for group in doc.groups: + if 'Члены' in group.position_in_group: + output_text += f'{group.group_name}. Должность внутри группы: {group.position_in_group.replace("Члены", "Член")}\n' + else: + output_text += f'{group.group_name}. Должность внутри группы: {group.position_in_group}\n' + output_text += f'\\\n\n' + + else: + if isinstance(answer_chunks[0], FilterChunks): + doc = answer_chunks[0] + output_text += f'Документ: [1]\n' + if doc.title != 'unknown': + output_text += f'Название документа: {doc.title}\n' + for chunk in doc.chunks: + if len(chunk.other_info): + for i in chunk.other_info: + output_text += f'{i}' + else: + output_text += f'{chunk.text_answer}' + output_text += '\n\n' + else: + doc = answer_chunks[0] + output_text += f'Информация о сотруднике {doc.person_name}\n' + if doc.organizatinal_structure is not None: + for organizatinal_structure in doc.organizatinal_structure: + output_text += ( + f'Должность: {organizatinal_structure.position}\n' + ) + if organizatinal_structure.leads is not None: + output_text += f'Руководит следующими сотрудниками:\n' + for lead in organizatinal_structure.leads: + if lead.person != "undefined": + output_text += f'{lead.person}\n' + if ( + organizatinal_structure.subordinates.person_name + != "undefined" + ): + output_text += f'Руководителем {doc.person_name} является {organizatinal_structure.subordinates.person_name}\n' + + if doc.business_processes is not None: + if len(doc.business_processes) >= 2: + output_text += f'Отвечает за Бизнес процессы:\n' + else: + output_text += f'Отвечает за Бизнес процесс: ' + for process in doc.business_processes: + output_text += f'{process.processes_name}\n' + if doc.business_curator is not None: + output_text += 'Является Бизнес-куратором (РОКС НН):\n' + for curator in doc.business_curator: + output_text += f'{curator.company_name}' + if doc.groups is not None: + if len(doc.groups) >= 2: + output_text += 'Входит в состав групп:\n' + else: + output_text += 'Входит в состав группы:\n' + for group in doc.groups: + if 'Члены' in group.position_in_group: + output_text += f'{group.group_name}. Должность внутри группы: {group.position_in_group.replace("Члены", "Член")}\n' + else: + output_text += f'{group.group_name}. Должность внутри группы: {group.position_in_group}\n' + output_text += f'\\\n\n' + + return output_text diff --git a/components/nmd/metadata_manager.py b/components/nmd/metadata_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..d0d909dea74256222b5b7fd160f4e0f4962363fc --- /dev/null +++ b/components/nmd/metadata_manager.py @@ -0,0 +1,255 @@ +from typing import List, Tuple, Optional + +import pandas as pd + + +class MetadataManager: + def __init__(self, df: pd.DataFrame, logger): + self.logger = logger + self.df = df + self.df.drop('Embedding', axis=1, inplace=True) + self.df = self.df.where(pd.notna(self.df), 'unknown') + + @staticmethod + def __search_sub_level(df: pd.DataFrame, header_text: Optional[str] = None) -> List: + """ + + Args: + df: + + Returns: + + """ + paragraphs = [] + if header_text is None: + header_text = df.iloc[0]['Text'] + + for ind, (_, row) in enumerate(df.iterrows()): + text = row['Text'] + if ind == 0: + text = text.replace(f'{header_text}', f'{header_text}\n') + else: + text = text.replace(f'{header_text}', '') + '\n' + paragraphs.append(text) + return paragraphs + + @staticmethod + def __check_duplicates(df: pd.DataFrame, ind: int) -> pd.DataFrame: + if df.loc[ind]['Duplicate'] is not None: + return df[df['Duplicate'] == df.loc[ind]['Duplicate']] + else: + return df[df['Duplicate'].isna()] + + @staticmethod + def __check_appendix_duplicates(df: pd.DataFrame, ind: int) -> pd.DataFrame: + if df.loc[ind]['DuplicateAppendix'] is not None: + return df[df['DuplicateAppendix'] == df.loc[ind]['DuplicateAppendix']] + else: + return df[df['DuplicateAppendix'].isna()] + + def _paragraph_appendix_content(self, df, pattern: str, ind: int, shape: int) -> Tuple[List, int]: + """ + Функция возвращает контент параграфа. Если в параграфе были подпункты через "-" или буквы "а, б" + Args: + df: DataFrame + pattern: Паттерн поиска. + ind: Индекс строки в DataFrame. + shape: Размер DataFrame при котором будет возвращаться пустой список. + + Returns: + Возвращает список подразделов. + Examples: + 3.1. Параграф: + 1) - Содержание 1; + 2) - Содержание 2; + 3) - Содержание 3; + """ + df = df[(df['PargaraphAppendix'].str.match(pattern, na=False)) | (df.index == ind)] + df = self.__check_appendix_duplicates(df, ind) + + if df.shape[0] <= shape: + return [], None + + start_index_paragraph = df.index[0] + paragraphs = self.__search_sub_level(df) + return paragraphs, start_index_paragraph + + def _paragraph_content(self, df, pattern: str, ind: int, shape: int) -> Tuple[List, int]: + """ + Функция возвращает контент параграфа. Если в параграфе были подпункты через "-" или буквы "а, б" + Args: + df: DataFrame + pattern: Паттерн поиска. + ind: Индекс строки в DataFrame. + shape: Размер DataFrame при котором будет возвращаться пустой список. + + Returns: + Возвращает список подразделов. + Examples: + 3.1. Параграф: + 1) - Содержание 1; + 2) - Содержание 2; + 3) - Содержание 3; + """ + df = df[ + (df['Pargaraph'].str.match(pattern, na=False)) & # Проверка, соответствуют ли значения паттерну + (df['Duplicate'] == df.loc[ind]['Duplicate']) | # Оставить разделы только принадлежащие одному дубликату + (df.index == ind)] # Оставить значение, которое нашел векторный поиск + # df = self.__check_duplicates(df, ind) + + if df.shape[0] <= shape: + return [], None + + start_index_paragraph = df.index[0] + paragraphs = self.__search_sub_level(df) + return paragraphs, start_index_paragraph + + def _paragraph_content2(self, df, pattern: str, ind: int, shape: int) -> Tuple[List, int]: + """ + Функция возвращает контент параграфа. Если в параграфе были подпункты через "-" или буквы "а, б" + Args: + df: DataFrame + pattern: Паттерн поиска. + ind: Индекс строки в DataFrame. + shape: Размер DataFrame при котором будет возвращаться пустой список. + + Returns: + Возвращает список подразделов. + Examples: + 3.1. Параграф: + 1) - Содержание 1; + 2) - Содержание 2; + 3) - Содержание 3; + """ + df = df[df['Pargaraph'].str.match(pattern, na=False)] + if df.shape[0] <= shape: + return [], None + # df = self.__check_duplicates(df, ind) + # if df.shape[0] <= shape: + # return [], None + start_index_paragraph = df.index[0] + paragraphs = self.__search_sub_level(df) + return paragraphs, start_index_paragraph + + @staticmethod + def _first_unknown_index(df): + indexes = list(df[df['PartLevel1'].isin(['unknown'])].index) + if len(indexes) > 0: + return df.loc[indexes[-1]]['Text'] + else: + return None + + def _search_other_info(self, ind, doc_number): + + df = self.df[self.df['DocNumber'] == doc_number] + start_index_paragraph = df.loc[ind]['Index'] - 1 + if df.loc[ind]['Table'] != 'unknown': + return df.loc[ind]['Text'], ind + + if df.loc[ind]['PartLevel1'] != 'unknown': + if 'Table' in str(self.df.iloc[ind]['PartLevel1']): + return [], ind + + if df.loc[ind]['Appendix'] != 'unknown': + df = df[df['Appendix'] == self.df.iloc[ind]['Appendix']] + if df.loc[ind]['LevelParagraphAppendix'] == 'unknown' and df.loc[ind]['PargaraphAppendix'] == 'unknown': + # pattern = r'\d+\.?$' + # df = df[(df['PargaraphAppendix'].str.match(pattern, na=False)) | (df.index == ind)] + # df = df[(df['LevelParagraphAppendix'] == 'Level0') | (df.index == ind)] + df = df.loc[ind:ind + 7] + start_index_paragraph = df.index[0] + paragraph = self.__search_sub_level(df) + elif df.loc[ind]['PargaraphAppendix'] != 'unknown': + pattern = df.loc[ind]["PargaraphAppendix"].replace(".", r"\.") + pattern = f'^{pattern}?\\d?.?$' + if df[df['PargaraphAppendix'].str.match(pattern, na=False)].shape[0] == 1: + pattern = df.loc[ind]["PargaraphAppendix"].replace(".", r"\.") + pattern = pattern.split('.') + pattern = [elem for elem in pattern if elem] + if len(pattern) == 1: + pattern = '.'.join(pattern) + pattern = f'^{pattern}.?\\d?.?$' + else: + pattern = '.'.join(pattern[:-1]) + pattern = f'^{pattern}.\\d.?$' + df = df[df['PargaraphAppendix'].str.match(pattern, na=False)] + start_index_paragraph = df.index[0] + paragraph = self.__search_sub_level(df) + else: + paragraph = self.df.iloc[int(ind - 10):ind + 10]['Text'].values + start_index_paragraph = df.index[0] + return ' '.join(paragraph), start_index_paragraph + else: + if df.loc[ind]['Pargaraph'] == 'unknown': + header_text = self._first_unknown_index(df) + df = df.loc[int(ind - 2):ind + 2] + paragraph = self.__search_sub_level(df, header_text) + # Связан с документами без пунктов поэтому передается несколько параграфов сверху и снизу + else: + pattern = df.loc[ind]["Pargaraph"].replace(".", r"\.") + # Изет под пункты внутри пункта + paragraph, start_index_paragraph = self._paragraph_content(df, fr'^{pattern}?$', ind, 2) + if len(paragraph) == 0: + pattern = f'{pattern}\\d?.?\\d?\\d?.?$' + paragraph, start_index_paragraph = self._paragraph_content2(df, pattern, ind, 0) + if len(paragraph) == 0 and df.loc[ind]['LevelParagraph'] != '0': + pattern = df.loc[ind]["Pargaraph"].split('.') + pattern = [elem for elem in pattern if elem] + pattern = '.'.join(pattern[:-1]) + pattern = f'^{pattern}\\.\\d\\d?.?$' + paragraph, start_index_paragraph = self._paragraph_content(df, pattern, ind, 0) + elif len(paragraph) == 0 and df.loc[ind]['LevelParagraph'] == '0': + pattern = df.loc[ind]["Pargaraph"].replace(".", r"\.") + if '.' not in pattern: + pattern = pattern + '\.' + pattern = f'^{pattern}\\d.?\\d?.?$' + paragraph, start_index_paragraph = self._paragraph_content(df, pattern, ind, 0) + + return ' '.join(paragraph), start_index_paragraph + + @staticmethod + def filter_answer(answer): + flip_answer = [] + new_answer = {} + count = 0 + for key in answer: + if answer[key]['start_index_paragraph'] not in flip_answer: + flip_answer.append(answer[key]['start_index_paragraph']) + new_answer[count] = answer[key] + count += 1 + return new_answer + + def _clear_doc_name(self, ind): + split_doc_name = self.df.iloc[ind]['DocName'].split('_') + return ' '.join(split_doc_name[1:]).replace('.txt', '').replace('.json', '').replace('.DOCX', '').replace( + '.DOC', '').replace('tables', '') + + def search(self, indexes: List) -> dict: + """ + Метод ищет ответы на запрос + Args: + indexes: Список индексов. + + Returns: + Возвращает словарь с ответами и информацией об ответах. + """ + answers = {} + for i, ind in enumerate(indexes): + answers[i] = {} + doc_number = self.df.iloc[ind]['DocNumber'] + answers[i]['id'] = doc_number + answers[i][f'index_answer'] = int(ind) + answers[i][f'doc_name'] = self._clear_doc_name(ind) + answers[i][f'title'] = self.df.iloc[ind]['Title'] + answers[i][f'text_answer'] = self.df.iloc[ind]['Text'] + + try: + other_info, start_index_paragraph = self._search_other_info(ind, doc_number) + except KeyError: + other_info, start_index_paragraph = self.df.iloc[ind]['Text'], ind + self.logger.info('Ошибка в индексе, проверьте БД!') + if len(other_info) == 0: + other_info, start_index_paragraph = self.df.iloc[ind]['Text'], ind + answers[i][f'other_info'] = [other_info] + answers[i][f'start_index_paragraph'] = int(start_index_paragraph) + return self.filter_answer(answers) diff --git a/components/nmd/query_classification.py b/components/nmd/query_classification.py new file mode 100644 index 0000000000000000000000000000000000000000..8513a05150a71c4839bad395f7a3553dd1b58a0e --- /dev/null +++ b/components/nmd/query_classification.py @@ -0,0 +1,79 @@ +import os +import re +from logging import Logger +from typing import Dict, List, Optional, Tuple + +from openai import OpenAI + +from common.configuration import LLMConfiguration + + +class QueryClassification: + def __init__(self, config: LLMConfiguration, prompt: str, logger: Logger): + self.config = config + self.logger = logger + self.prompt = prompt + self.pattern = r'\[\d+\]' + + # Initialize OpenAI client + if self.config.base_url is not None: + self.client = OpenAI( + base_url=self.config.base_url, + api_key=os.getenv(self.config.api_key_env) + ) + else: + self.client = None + + def query_classification(self, query: str) -> Tuple[str, Optional[Dict], Optional[List]]: + """ + Classify the query using LLM + + Args: + query: User query to classify + + Returns: + Tuple containing query type, optional metadata and optional list + """ + self.logger.info('Query Classification') + if self.client is None: + return '[3]', None, None + for i in range(5): + try: + response = self.client.chat.completions.create( + model=self.config.model, + messages=[ + {"role": "system", "content": self.prompt}, + {"role": "user", "content": query} + ], + temperature=self.config.temperature, + top_p=self.config.top_p, + frequency_penalty=self.config.frequency_penalty, + presence_penalty=self.config.presence_penalty, + seed=self.config.seed + ) + + answer_llm = response.choices[0].message.content + self.logger.info(f'Answer LLM {answer_llm}') + + # Process the response + if re.search('%%', answer_llm): + index = re.search('%%', answer_llm).span()[1] + answer_llm = answer_llm[index:] + if re.search('Конец ответа', answer_llm): + index = re.search('Конец ответа', answer_llm).span()[1] + answer_llm = answer_llm[:index] + + # Extract query type + query_type = re.findall(self.pattern, answer_llm) + if query_type: + query_type = query_type[0] + else: + query_type = '[3]' + + return query_type, None, None + + except Exception as e: + self.logger.error(f"Attempt {i+1} failed: {str(e)}") + if i == 4: + self.logger.error("All attempts failed") + return '[3]', None, None diff --git a/components/nmd/rancker.py b/components/nmd/rancker.py new file mode 100644 index 0000000000000000000000000000000000000000..79a589ef0eee887f5f9c920eb229544555ae4b58 --- /dev/null +++ b/components/nmd/rancker.py @@ -0,0 +1,32 @@ +import pandas as pd + +from common.configuration import Configuration + + +class DocumentRanking: + + def __init__(self, df: pd.DataFrame, config: Configuration): + self.df = df + self.config = config + self.alpha = config.db_config.ranker.alpha + self.beta = config.db_config.ranker.beta + + def doc_ranking(self, query_embedding, scores, indexes): + title_embeddings = self.df.iloc[indexes]['TitleEmbedding'].to_list() + norms = [] + for emb in title_embeddings: + d = emb - query_embedding + norm = d.dot(d) + norms.append(norm) + + new_score = [] + texts = self.df.iloc[indexes]['Text'].to_list() + for ind, text in enumerate(texts): + new_score.append(scores[ind] * len(text) ** self.beta + self.alpha * norms[ind]) + + metric_df = pd.DataFrame() + metric_df['NewScores'] = new_score + metric_df['Indexes'] = indexes + metric_df.sort_values(by=['NewScores'], inplace=True) + new_indexes = metric_df['Indexes'].to_list()[:self.config.db_config.search.vector_search.k_neighbors] + return new_indexes diff --git a/components/parser/README.md b/components/parser/README.md new file mode 100644 index 0000000000000000000000000000000000000000..be698b67114270ead091cf07392396ab456a7cef --- /dev/null +++ b/components/parser/README.md @@ -0,0 +1,105 @@ +# Pipeline Module + +> ВАЖНО!!! README.md сгенерировано автоматически, поэтому может содержать неточности. + +Модуль реализует пайплайн для обработки XML документов и создания структурированного датасета. Пайплайн включает несколько последовательных этапов обработки, от парсинга XML до создания векторизованного датасета. + +## Основные этапы обработки + +### 1. Парсинг XML файлов +- Чтение XML файлов из указанной директории +- Извлечение текстового и табличного контента +- Сохранение метаданных документов + +### 2. Обработка аббревиатур +- Извлечение аббревиатур из текста документов +- Объединение с предварительно подготовленными аббревиатурами +- Применение аббревиатур к текстовому и табличному контенту +- Сохранение списка обнаруженных аббревиатур + +### 3. Извлечение иерархической структуры +- Парсинг структуры текстового контента +- Парсинг структуры табличного контента +- Создание иерархического представления документов + +### 4. Создание датасета +- Формирование структурированного датасета +- Векторизация текстов +- Сохранение результатов + +## Использование + +```python +from components.embedding_extraction import EmbeddingExtractor +from components.parser.pipeline import DatasetCreationPipeline +from components.parser.abbreviations.abbreviation import Abbreviation + +# Инициализация пайплайна +pipeline = DatasetCreationPipeline( + dataset_id="my_dataset", + vectorizer=EmbeddingExtractor(), + prepared_abbreviations=[], # список предварительно подготовленных аббревиатур + xml_ids=["doc1", "doc2"], # список идентификаторов XML файлов + save_intermediate_files=True # сохранять ли промежуточные файлы +) + +# Запуск пайплайна +dataset = pipeline.run() +``` + +## Структура выходных данных + +### Основные файлы +- `dataset.csv` - финальный датасет с векторизованными текстами +- `abbreviations.csv` - извлеченные аббревиатуры +- `xml_info.csv` - метаданные XML документов + +### Промежуточные файлы (опционально) +- `txt/*.txt` - извлеченный текстовый контент +- `txt_abbr/*.txt` - текстовый контент после применения аббревиатур +- `jsons/*.json` - иерархическая структура документов + +## Параметры конфигурации + +### DatasetCreationPipeline +- `dataset_id: str` - идентификатор создаваемого датасета +- `vectorizer: EmbeddingExtractor` - векторизатор для создания эмбеддингов +- `prepared_abbreviations: list[Abbreviation]` - предварительно подготовленные аббревиатуры +- `xml_ids: list[str]` - список идентификаторов XML файлов для обработки +- `save_intermediate_files: bool` - сохранять ли промежуточные файлы + +## Зависимости + +### Внутренние компоненты +- `components.embedding_extraction.EmbeddingExtractor` +- `components.parser.abbreviations.AbbreviationExtractor` +- `components.parser.features.HierarchyParser` +- `components.parser.features.DatasetCreator` +- `components.parser.xml.XMLParser` + +### Внешние библиотеки +- pandas +- numpy +- pathlib + +## Структура директорий + +``` +data/ +└── regulation_datasets/ + └── {dataset_id}/ + ├── abbreviations.csv + ├── xml_info.csv + ├── dataset.csv + ├── embeddings.pt + ├── txt/ # (опционально) + ├── txt_abbr/ # (опционально) + └── jsons/ # (опционально) +``` + +## Примечания + +- Все промежуточные файлы сохраняются только если установлен флаг `save_intermediate_files=True` +- Векторизация выполняется после создания датасета +- Аббревиатуры применяются как к текстовому, так и к табличному контенту +- Иерархическая структура извлекается отдельно для текста и таблиц \ No newline at end of file diff --git a/components/parser/abbreviations/README.md b/components/parser/abbreviations/README.md new file mode 100644 index 0000000000000000000000000000000000000000..1d0ab4927d4c2c95e517118b3a3ed7684a4fcc16 --- /dev/null +++ b/components/parser/abbreviations/README.md @@ -0,0 +1,119 @@ +# Экстрактор сокращений (Abbreviation Extractor) + +> ВАЖНО!!! README.md сгенерировано автоматически, поэтому может содержать неточности. + +Модуль для извлечения сокращений и их полных форм из текстовых документов. + +## Принцип работы + +Экстрактор ищет в тексте конструкции вида: +- "полная форма (далее - сокращение)" +- "полная форма (далее – сокращение)" +и подобные варианты. + +### Основные этапы обработки: + +1. **Разбиение на предложения** + - Текст разбивается на предложения с учетом специальных случаев + - Учитываются особые сокращения, после которых точка не является концом предложения + +2. **Поиск сокращений** + - В каждом предложении ищутся конструкции с маркером "далее" + - Извлекается короткая форма (сокращение) после маркера + - Определяется полная форма до маркера + +3. **Обработка сокращений** + - Поддерживается два типа сокращений: + - Однословные (например, "*БЖВРК*") + - Многословные (например, "Мы великая нация великих обезьян (далее - *нация обезьян*)") + - Для каждого сокращения определяется его полная форма + +4. **Лемматизация** + - Используется библиотека Natasha для лемматизации текста + - Помогает находить соответствия между полной и короткой формами + +## Использование + +```python +from components.parser.abbreviations.abbreviation_extractor import AbbreviationExtractor +from components.parser.xml.structures import ParsedXMLs + +# Создание экстрактора +extractor = AbbreviationExtractor() + +# Обработка XML-файлов +result = extractor.process_parsed_xmls(parsed_xmls) + +# Обработка одного файла +file_abbreviations = extractor.process_file(text, filename) + +# Извлечение сокращений из текста +abbreviations = extractor.extract_abbreviations_from_text(text) +``` + +## Структура результатов + +Результаты представляются в виде структур данных: +- `AllFilesAbbreviations` - коллекция сокращений из всех файлов +- `OneFileAbbreviations` - сокращения из одного файла +- `Abbreviation` - отдельное сокращение с полной и короткой формами + +## Особенности + +- Учитываются различные варианты разделителей между полной и короткой формами +- Поддерживается обработка специальных сокращений, не являющихся концом предложения +- Используется лемматизация для улучшения поиска соответствий +- Возможна обработка как одиночных файлов, так и наборов файлов + +# Обработка сокращений и аббревиатур + +Модуль `abbreviation.py` отвечает за обработку и нормализацию сокращений и аббревиатур в тексте. + +## Основные типы сокращений + +- `ABBREVIATION` - аббревиатуры (например, "ОКС НН") +- `SHORTENING` - сокращения (например, "Компания") +- `UNKNOWN` - неопределенный тип + +## Процесс обработки + +Класс `Abbreviation` выполняет следующие этапы обработки: + +1. **Определение типа сокращения** (`_define_abbreviation_type`): + - Проверяет, является ли строка аббревиатурой (содержит более одной заглавной буквы в каждом слове) + - Проверяет, является ли строка сокращением (одно слово, начинающееся с заглавной буквы) + +2. **Очистка префиксов** (`_remove_prefix`): + - Удаляет такие префиксы как "далее", различные виды тире + - Убирает лишние пробелы + +3. **Очистка от мусора** (`_remove_trash`): + - Удаляет такие подстроки как "ПАО", "ОАО", "№", "(" + - Обрезает строку с начала до первого вхождения "мусорной" подстроки + +4. **Специальная обработка для аббревиатур** (`_process_abbreviation`): + - Извлекает заглавные буквы из короткой формы + - Проверяет соответствие заглавных букв началам слов в полной форме + - Обрезает полную форму до релевантной части + +5. **Специальная обработка для сокращений** (`_process_shortening`): + - Применяет стемминг (с помощью алгоритма Портера) к короткой форме + - Обрезает полную форму до релевантной части + +## Валидация + +- Проверяет длину полной формы (должна быть меньше MAX_LENGTH) +- Проверяет, что полная форма длиннее короткой +- Проверяет отсутствие полной формы в черном списке (BLACKLIST) +- Для аббревиатур проверяет соответствие заглавных букв началам слов +- Для сокращений проверяет корректность регистра букв и отсутствие специальных случаев + +Если какая-либо проверка не проходит, тип сокращения устанавливается как `UNKNOWN`. + + +# Применение сокращений и аббревиатур + +Класс `Abbreviation` имеет метод `apply`, который принимает текст и возвращает текст с примененными сокращениями и аббревиатурами. + +Класс + diff --git a/components/parser/abbreviations/__init__.py b/components/parser/abbreviations/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c7312f90aab2d52dae6b04a2e4d94f267599d703 --- /dev/null +++ b/components/parser/abbreviations/__init__.py @@ -0,0 +1,9 @@ +from .abbreviation import Abbreviation +from .abbreviation_extractor import AbbreviationExtractor +from .structures import AbbreviationsCollection + +__all__ = [ + "AbbreviationExtractor", + "Abbreviation", + "AbbreviationsCollection", +] diff --git a/components/parser/abbreviations/abbreviation.py b/components/parser/abbreviations/abbreviation.py new file mode 100644 index 0000000000000000000000000000000000000000..33e2c369b566b5a154c95b2cce997fce15b0c4cb --- /dev/null +++ b/components/parser/abbreviations/abbreviation.py @@ -0,0 +1,328 @@ +import re +from dataclasses import dataclass +from enum import Enum + +from components.parser.abbreviations.constants import ( + ABBREVIATION_CLEANUP_REPLACEMENTS, + BLACKLIST, + DASH_PATTERN, + MAX_LENGTH, + PREFIX_PARTS_TO_REMOVE, + REMOVING_SUBSTRINGS, +) +from components.parser.abbreviations.porter import Porter + + +class AbbreviationType(str, Enum): + ABBREVIATION = 'abbreviation' + SHORTENING = 'shortening' + UNKNOWN = 'unknown' + + +@dataclass +class Abbreviation: + short_form: str + full_form: str + abbreviation_type: AbbreviationType = AbbreviationType.UNKNOWN + + _processed: bool = False + document_id: int | None = None + + def process(self) -> 'Abbreviation': + """ + Производит пост-обработку сокращения и полной формы. + - Определяет тип сокращения. + - Удаляет префикс из короткой формы и мусор из полной формы. + - В зависимости от типа сокращения адаптирует его под нужный вид. + """ + if self._processed: + return + + self._define_abbreviation_type() + + self.short_form = self._remove_prefix(self.short_form) + self.full_form = self._remove_trash(self.full_form) + + if self._abbreviation_type == AbbreviationType.SHORTENING: + self._process_shortening() + elif self._abbreviation_type == AbbreviationType.ABBREVIATION: + self._process_abbreviation() + + self._processed = True + + return self + + def apply(self, text: str) -> str: + """ + Применяет аббревиатуру к тексту. + + Args: + text (str): Текст для обработки. + + Returns: + str: Обработанный текст. + """ + if self._abbreviation_type == AbbreviationType.UNKNOWN: + return text + + if self._abbreviation_type == AbbreviationType.SHORTENING: + return self._apply_shortening(text) + + elif self._abbreviation_type == AbbreviationType.ABBREVIATION: + return self._apply_abbreviation(text) + + def _apply_shortening(self, text: str) -> str: + """ + Применяет сокращение к тексту. + + Args: + text (str): Текст для обработки. + + Returns: + str: Обработанный текст. + """ + matches = list(re.finditer(self.short_form, text)) + for i in range(len(matches) - 1, 1, -1): + m = matches[i] + pos1 = m.start() + m2 = re.match(r'[A-Za-zА-Яа-я]+', text[pos1:]) + pos2 = pos1 + m2.end() + explanation = self.full_form + m3 = re.match(r'[A-Za-zА-Яа-я]+', explanation) + explanation = explanation[m3.end() :] + text = text[:pos2] + explanation + text[pos2:] + return text + + def _apply_abbreviation(self, text: str) -> str: + """ + Применяет аббревиатуру к тексту. + + Args: + text (str): Текст для обработки. + + Returns: + str: Обработанный текст. + """ + matches = list(re.finditer(self.short_form, text)) + for i in range(len(matches) - 1, 0, -1): + m = matches[i] + text = f'{text[: m.start()]}{self.short_form} ({self.full_form}){text[m.end():]}' + return text + + def _define_abbreviation_type(self) -> None: + """ + Определяет тип сокращения. + """ + if self._check_abbreviation(self.full_form): + self._abbreviation_type = AbbreviationType.ABBREVIATION + elif self._check_shortening(self.full_form): + self._abbreviation_type = AbbreviationType.SHORTENING + else: + self._abbreviation_type = AbbreviationType.UNKNOWN + + def _process_shortening(self) -> None: + """ + Обрабатывает сокращение. + """ + key = Porter.stem(self.short_form) + pos = self.full_form.lower().rfind(key.lower()) + if pos != -1: + self.full_form = self.full_form[pos:] + self.short_form = key + else: + self.abbreviation_type = AbbreviationType.UNKNOWN + + def _process_abbreviation(self) -> None: + """ + Обрабатывает аббревиатуру. + """ + uppercase_letters = re.sub('[a-zа-я, ]', '', self.short_form) + processed_full_form = self._remove_trash_when_abbreviation(self.full_form) + words = processed_full_form.split() + uppercase_letters = uppercase_letters[::-1] + words = words[::-1] + + if (len(words) <= len(uppercase_letters)) or ('ОКС НН' not in self.short_form): + self.abbreviation_type = AbbreviationType.UNKNOWN + return + + match = self._check_abbreviation_matches_words(uppercase_letters, words) + if match: + self._process_matched_abbreviation(uppercase_letters, words) + else: + self._process_mismatched_abbreviation() + + def _process_matched_abbreviation( + self, + uppercase_letters: str, + words: list[str], + ) -> None: + """ + Обрабатывает аббревиатуру, которая совпадает с первыми буквами полной формы. + + Args: + uppercase_letters (str): Заглавные буквы из сокращения. + words (list[str]): Список слов, которые составляют аббревиатуру. + """ + pos = len(self.full_form) + for i in range(len(uppercase_letters)): + pos = self.full_form.rfind(words[i], 0, pos) + + if pos != -1: + self.full_form = self.full_form[pos:] + + else: + self.abbreviation_type = AbbreviationType.UNKNOWN + + def _process_mismatched_abbreviation(self) -> None: + """ + Обрабатывает аббревиатуру, которая не совпадает с первыми буквами полной формы. + """ + first_letter = self.short_form[0] + pos = self.full_form.rfind(first_letter) + if pos != -1: + self.full_form = self.full_form[pos:] + first_letter = self.full_form[0] + second_letter = self.full_form[1] + + if ( + ('A' < first_letter < 'Z' or 'А' < first_letter < 'Я') + and ('a' < second_letter < 'z' or 'а' < second_letter < 'я') + and len(self.full_form) < MAX_LENGTH + and len(self.full_form) > len(self.short_form) + and self.full_form not in BLACKLIST + and '_' not in self.full_form + ): + return + + self.abbreviation_type = AbbreviationType.UNKNOWN + + def _check_abbreviation_matches_words( + self, + uppercase_letters: str, + words: list[str], + ) -> bool: + """ + Проверяет, соответствует ли короткая форма аббревиатуре. + + Args: + uppercase_letters (str): Заглавные буквы из сокращения. + words (list[str]): Список слов, которые составляют аббревиатуру. + + Returns: + bool: True, если аббревиатура соответствует, False в противном случае. + """ + for j in range(len(uppercase_letters)): + c1 = uppercase_letters[j].lower() + c2 = words[j][0].lower() + if c1 != c2: + return False + + return True + + @classmethod + def _check_abbreviation(cls, full_form: str) -> bool: + """ + Проверяет, является ли строка аббревиатурой. + + Args: + full_form (str): Строка для проверки. + + Returns: + bool: True, если строка является аббревиатурой, False в противном случае. + """ + s = cls._remove_prefix(full_form) + words = s.split() + + for word in words: + n = cls._count_uppercase_letters(word) + if (n <= 1) and (word != 'и'): + return False + + return True + + @classmethod + def _check_shortening(cls, full_form: str) -> bool: + """ + Проверяет, является ли строка сокращением. + + Args: + full_form (str): Строка для проверки. + + Returns: + bool: True, если строка является сокращением, False в противном случае. + """ + s = cls._remove_prefix(full_form) + words = s.split() + + if len(words) != 1: + return False + + word = words[0] + if word[0].isupper() and word[1:].islower() and ('Компания' not in word): + return True + + return False + + @staticmethod + def _remove_prefix(s: str) -> str: + """ + Удаляет из строки префиксы типа "далее - " и "далее – ". + + Args: + s (str): Строка для обработки. + + Returns: + str: Обработанная строка. + """ + for prefix_part in PREFIX_PARTS_TO_REMOVE: + s = s.replace(prefix_part, '') + return s.strip() + + @staticmethod + def _remove_trash(s: str) -> str: + """ + Удаляет из строки такие подстроки, как "ПАО", "ОАО", "№", "(". + + Args: + s (str): Строка для обработки. + + Returns: + str: Обработанная строка. + """ + for substring in REMOVING_SUBSTRINGS: + pos = s.find(substring) + if pos != -1: + s = s[:pos] + return s + + @staticmethod + def _remove_trash_when_abbreviation(s: str) -> str: + """ + Удаляет из строки такие подстроки, как " и ", " или ", ", ", " ГО". + Заменяет дефисы и тире на пробел. + Это необходимо для того, чтобы правильно сопоставить аббревиатуру с полной формой. + + Args: + s (str): Строка для обработки. + + Returns: + str: Обработанная строка. + """ + for old, new in ABBREVIATION_CLEANUP_REPLACEMENTS.items(): + s = s.replace(old, new) + s = re.sub(DASH_PATTERN, ' ', s) + return s + + @staticmethod + def _count_uppercase_letters(s: str) -> int: + """ + Считает количество заглавных букв в строке. + + Args: + s (str): Строка для обработки. + + Returns: + int: Количество заглавных букв. + """ + return len(re.findall(r'[A-Z,А-Я]', s)) diff --git a/components/parser/abbreviations/abbreviation_extractor.py b/components/parser/abbreviations/abbreviation_extractor.py new file mode 100644 index 0000000000000000000000000000000000000000..111ea1eae73984af4157e2780ad0950e0bdb6743 --- /dev/null +++ b/components/parser/abbreviations/abbreviation_extractor.py @@ -0,0 +1,336 @@ +import re + +from natasha import Doc, MorphVocab, NewsEmbedding, NewsMorphTagger, Segmenter + + +from .constants import ( + ABBREVIATION_RE, + CLOSE_BRACKET_RE, + FIRST_CHARS_SET, + NEXT_MARKER_RE, + NON_SENTENCE_ENDINGS, + SECOND_CHARS_SET, + UPPERCASE_LETTER_RE, +) +from .structures import Abbreviation + + +class AbbreviationExtractor: + def __init__(self): + """ + Инициализация экстрактора сокращений. + + Создает необходимые компоненты для лемматизации и компилирует регулярные выражения. + """ + # Инициализация компонентов Natasha для лемматизации + self.segmenter = Segmenter() + self.morph_tagger = NewsMorphTagger(NewsEmbedding()) + self.morph_vocab = MorphVocab() + + # Компиляция регулярных выражений + self.next_re = re.compile(NEXT_MARKER_RE, re.IGNORECASE) + self.abbreviation_re = re.compile(ABBREVIATION_RE) + self.uppercase_letter_re = re.compile(UPPERCASE_LETTER_RE) + self.close_bracket_re = re.compile(CLOSE_BRACKET_RE) + + self.delimiters = [ + f'{char1} {char2} '.format(char1, char2) + for char1 in FIRST_CHARS_SET + for char2 in SECOND_CHARS_SET + ] + + def extract_abbreviations_from_text( + self, + text: str, + ) -> list[Abbreviation]: + """ + Извлечение всех сокращений из текста. + + Args: + text: Текст для обработки + + Returns: + list[Abbreviation]: Список найденных сокращений + """ + sentences = self._extract_sentences_with_abbreviations(text) + + abbreviations = [self._process_one_sentence(sentence) for sentence in sentences] + abbreviations = sum(abbreviations, []) # делаем список одномерным + abbreviations = [abbreviation.process() for abbreviation in abbreviations] + + return abbreviations + + def _process_one_sentence(self, sentence: str) -> list[Abbreviation]: + """ + Обработка одного предложения для извлечения сокращений. + + Args: + sentence: Текст для обработки + + Returns: + list[Abbreviation]: Список найденных сокращений + """ + search_iter = self.next_re.finditer(sentence) + prev_index = 0 + + abbreviations = [] + + for match in search_iter: + abbreviation, prev_index = self._process_match(sentence, match, prev_index) + if abbreviation is not None: + abbreviations.append(abbreviation) + + return abbreviations + + def _process_match( + self, + sentence: str, + match: re.Match, + prev_index: int, + ) -> tuple[Abbreviation | None, int]: + """ + Обработка одного совпадения с конструкцией "далее - {short_form}" для извлечения сокращений. + + Args: + sentence: Текст для обработки + match: Совпадение для обработки + prev_index: Предыдущий индекс + + Returns: + tuple[Abbreviation | None, int]: Найденное сокращение (None, если нет сокращения) и следующий индекс + """ + start, end = match.start(), match.end() + text = sentence[start:] + + index_close_parenthesis = self._get_close_parenthesis_index(text) + index_point = self._get_point_index(text, start) + + prev_index += index_point + short_word = text[end : start + index_close_parenthesis].strip() + + if len(short_word.split()) < 2: + abbreviation = self._process_match_for_word( + short_word, text, start, end, prev_index + ) + + else: + abbreviation = self._process_match_for_phrase( + short_word, text, start, end, prev_index + ) + + prev_index = start + index_close_parenthesis + 1 + + return abbreviation, prev_index + + def _get_close_parenthesis_index(self, text: str) -> int: + """ + Получение индекса закрывающей скобки в тексте. + + Args: + text: Текст для обработки + + Returns: + int: Индекс закрывающей скобки или 0, если не найдено + """ + result = self.close_bracket_re.search(text) + if result is None: + return 0 + return result.start() + + def _get_point_index(self, text: str, start_index: int) -> int: + """ + Получение индекса точки в тексте. + + Args: + text: Текст для обработки + start_index: Индекс начала поиска + + Returns: + int: Индекс точки или 0, если не найдено + """ + result = text.rfind('.', 0, start_index - 1) + if result == -1: + return 0 + return result + + def _process_match_for_word( + self, + short_word: str, + text: str, + start_next_re_index: int, + end_next_re_index: int, + prev_index: int, + ) -> Abbreviation | None: + """ + Обработка сокращения, состоящего из одного слова. + + Args: + short_word: Сокращение + text: Текст для обработки + start_next_re_index: Индекс начала следующего совпадения + end_next_re_index: Индекс конца следующего совпадения + prev_index: Предыдущий индекс + + Returns: + Abbreviation | None: Найденное сокращение или None, если нет сокращения + """ + if self.abbreviation_re.findall(text) or (short_word == 'ПДн'): + return None + + lemm_text = self._lemmatize_text(text[prev_index:start_next_re_index]) + lemm_short_word = self._lemmatize_text(short_word) + + search_word = re.search(lemm_short_word, lemm_text) + + if not search_word: + start_text_index = self._get_start_text_index( + text, + start_next_re_index, + prev_index, + ) + + if start_text_index is None: + return None + + full_text = text[prev_index + start_text_index : end_next_re_index] + + else: + index_word = search_word.span()[1] + space_index = text[prev_index:start_next_re_index].rfind(' ', 0, index_word) + if space_index == -1: + space_index = 0 + text = text[prev_index + space_index : start_next_re_index] + + full_text = text.replace(')', '').replace('(', '').replace('', '- ') + + return Abbreviation( + short_form=short_word, + full_form=full_text, + ) + + def _process_match_for_phrase( + self, + short_word: str, + text: str, + start_next_re_index: int, + end_next_re_index: int, + prev_index: int, + ) -> list[Abbreviation] | None: + """ + Обработка сокращения, состоящего из нескольких слов. + В действительности производится обработка первого слова сокращения, а затем вместо него подставляется полное сокращение. + + Args: + short_word: Сокращение + text: Текст для обработки + start_next_re_index: Индекс начала следующего совпадения + end_next_re_index: Индекс конца следующего совпадения + prev_index: Предыдущий индекс + + Returns: + list[Abbreviation] | None: Найденные сокращения или None, если нет сокращений + """ + first_short_word = short_word.split()[0] + result = self._process_match_for_word( + first_short_word, text, start_next_re_index, end_next_re_index, prev_index + ) + if result is None: + return None + return Abbreviation( + short_form=short_word, + full_form=result.full_form, + ) + + def _get_start_text_index( + self, + text: str, + start_next_re_index: int, + prev_index: int, + ) -> int | None: + """ + Получение индекса начала текста для поиска сокращения с учётом разделителей типа + "; - " + ": - " + ";  " + ": ‒ " и т.п. + + Args: + text: Текст для обработки + start_next_re_index: Индекс начала следующего совпадения + prev_index: Предыдущий индекс + + Returns: + int | None: Индекс начала текста или None, если не найдено + """ + if prev_index == 0: + return 0 + + for delimiter in self.delimiters: + result = re.search(delimiter, text[prev_index:start_next_re_index]) + if result is not None: + return result.span()[1] + + return None + + def _lemmatize_text(self, text: str) -> str: + """ + Лемматизация текста. + + Args: + text: Текст для лемматизации + + Returns: + str: Лемматизированный текст + """ + doc = Doc(text) + doc.segment(self.segmenter) + doc.tag_morph(self.morph_tagger) + + for token in doc.tokens: + token.lemmatize(self.morph_vocab) + + return ' '.join([token.lemma for token in doc.tokens]) + + def _extract_sentences_with_abbreviations(self, text: str) -> list[str]: + """ + Разбивает текст на предложения с учетом специальных сокращений. + + Точка после сокращений из NON_SENTENCE_ENDINGS не считается концом предложения. + + Args: + text: Текст для разбиения + + Returns: + list[str]: Список предложений + """ + text = text.replace('\n', ' ') + sentence_endings = re.finditer(r'\.\s+[А-Я]', text) + + sentences = [] + start = 0 + + for match in sentence_endings: + end = match.start() + 1 + + # Проверяем, не заканчивается ли предложение на специальное сокращение + preceding_text = text[start:end] + words = preceding_text.split() + + if words and any( + words[-1].rstrip('.').startswith(abbr) for abbr in NON_SENTENCE_ENDINGS + ): + continue + + sentence = text[start:end].strip() + sentences.append(sentence) + start = end + 1 + + # Добавляем последнее предложение + if start < len(text): + sentences.append(text[start:].strip()) + + return [ + sentence + for sentence in sentences + if self.next_re.search(sentence) is not None + ] diff --git a/components/parser/abbreviations/constants.py b/components/parser/abbreviations/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..f1ac45f5e055faa55f22dbdc2adf40b0bf2fafde --- /dev/null +++ b/components/parser/abbreviations/constants.py @@ -0,0 +1,54 @@ +# Регулярные выражения +NEXT_MARKER_RE = r'далее (--|־|᠆|‐|‑|‒|–|—|―|⸺|⸻|﹘|﹣|－|-|-)' +ABBREVIATION_RE = ( + r'\b[А-Я0-9]{1,}(?:\s?[А-Я0-9]{1,}|[:\-.]?[А-Я0-9]{1,}|[а-я]{1}[А-Я0-9]{1,})*\b' +) +UPPERCASE_LETTER_RE = r'[A-ZА-Я]' +CLOSE_BRACKET_RE = r'\)' + +# Сокращения, после которых точка не означает конец предложения +NON_SENTENCE_ENDINGS = ['г', 'д-р', 'т.е', 'и т.д', 'и т.п', 'и т.п.', 'ул', 'пр'] + +FIRST_CHARS_SET = {'.', ':', ';'} +SECOND_CHARS_SET = { + '‒', + '–', + '—', + '―', + '⸺', + '⸻', + '﹘', + '﹣', + '－', + '-', + '-', + '-', + '\uf0b7', + '', +} + +BLACKLIST = [ + 'Ненецкого муниципального района', + 'Изменение идентифицирующих', + 'Systems, Applications and Products in Data Processing Enterprise Resource Planning', + 'Российской Федерации, Уставом', + 'Собственника Объекта защиты', +] + +REMOVING_SUBSTRINGS = ['ПАО', 'ОАО', '№', '('] + +MAX_LENGTH = 100 + +# Strings to remove from abbreviations +PREFIX_PARTS_TO_REMOVE = ['далее', '–', '-'] + +# Strings to remove when processing abbreviations +ABBREVIATION_CLEANUP_REPLACEMENTS = { + ' и ': ' ', + ' или ': ' ', + ', ': ' ', + ' ГО': ' ', +} + +# Regex pattern for dashes/hyphens to be replaced with space +DASH_PATTERN = '(-|-|־|᠆|‐|‑|‒|–|—|―|⸺|⸻|﹘|﹣|－)' diff --git a/components/parser/abbreviations/porter.py b/components/parser/abbreviations/porter.py new file mode 100644 index 0000000000000000000000000000000000000000..13b5926f08e56e072f2205d2bbd2fd68cf7f3a95 --- /dev/null +++ b/components/parser/abbreviations/porter.py @@ -0,0 +1,63 @@ +import re + + +class Porter: + PERFECTIVEGROUND = re.compile( + u"((ив|ивши|ившись|ыв|ывши|ывшись)|((?<=[ая])(в|вши|вшись)))$" + ) + REFLEXIVE = re.compile(u"(с[яь])$") + ADJECTIVE = re.compile( + u"(ее|ие|ые|ое|ими|ыми|ей|ий|ый|ой|ем|им|ым|ом|его|ого|ему|ому|их|ых|ую|юю|ая|яя|ою|ею)$" + ) + PARTICIPLE = re.compile(u"((ивш|ывш|ующ)|((?<=[ая])(ем|нн|вш|ющ|щ)))$") + VERB = re.compile( + u"((ила|ыла|ена|ейте|уйте|ите|или|ыли|ей|уй|ил|ыл|им|ым|ен|ило|ыло|ено|ят|ует|уют|ит|ыт|ены|ить|ыть|ишь|ую|ю)|((?<=[ая])(ла|на|ете|йте|ли|й|л|ем|н|ло|но|ет|ют|ны|ть|ешь|нно)))$" + ) + NOUN = re.compile( + u"(а|ев|ов|ие|ье|е|иями|ями|ами|еи|ии|и|ией|ей|ой|ий|й|иям|ям|ием|ем|ам|ом|о|у|ах|иях|ях|ы|ь|ию|ью|ю|ия|ья|я)$" + ) + RVRE = re.compile(u"^(.*?[аеиоуыэюя])(.*)$") + DERIVATIONAL = re.compile(u".*[^аеиоуыэюя]+[аеиоуыэюя].*ость?$") + DER = re.compile(u"ость?$") + SUPERLATIVE = re.compile(u"(ейше|ейш)$") + I = re.compile(u"и$") + P = re.compile(u"ь$") + NN = re.compile(u"нн$") + + @staticmethod + def stem(word): + # word = word.lower() + word = word.replace(u'ё', u'е') + m = re.match(Porter.RVRE, word) + if m and m.groups(): + pre = m.group(1) + rv = m.group(2) + temp = Porter.PERFECTIVEGROUND.sub('', rv, 1) + if temp == rv: + rv = Porter.REFLEXIVE.sub('', rv, 1) + temp = Porter.ADJECTIVE.sub('', rv, 1) + if temp != rv: + rv = temp + rv = Porter.PARTICIPLE.sub('', rv, 1) + else: + temp = Porter.VERB.sub('', rv, 1) + if temp == rv: + rv = Porter.NOUN.sub('', rv, 1) + else: + rv = temp + else: + rv = temp + + rv = Porter.I.sub('', rv, 1) + + if re.match(Porter.DERIVATIONAL, rv): + rv = Porter.DER.sub('', rv, 1) + + temp = Porter.P.sub('', rv, 1) + if temp == rv: + rv = Porter.SUPERLATIVE.sub('', rv, 1) + rv = Porter.NN.sub(u'н', rv, 1) + else: + rv = temp + word = pre + rv + return word diff --git a/components/parser/abbreviations/structures.py b/components/parser/abbreviations/structures.py new file mode 100644 index 0000000000000000000000000000000000000000..b1a7f3022d226d707cd888168c0aba613e8a861f --- /dev/null +++ b/components/parser/abbreviations/structures.py @@ -0,0 +1,61 @@ +import logging +from dataclasses import dataclass + +import pandas as pd + +from components.parser.abbreviations.abbreviation import Abbreviation, AbbreviationType + +logger = logging.getLogger(__name__) + + +@dataclass +class AbbreviationsCollection: + items: list[Abbreviation] + + def to_pandas(self) -> pd.DataFrame: + """ + Преобразование всех сокращений в DataFrame. + + Returns: + pd.DataFrame: DataFrame с сокращениями + """ + logger.debug(f"Items: {self.items}") + all_data = [ + { + 'ShortWord': abbr.short_form, + 'LongText': abbr.full_form, + 'AbbreviationType': abbr.abbreviation_type, + 'DocumentId': abbr.document_id, + } + for abbr in self.items + if abbr.abbreviation_type != AbbreviationType.UNKNOWN + ] + logger.info(f'Approved abbreviations: {len(all_data)}') + logger.info(f'Rejected abbreviations: {len(self.items) - len(all_data)}') + return pd.DataFrame(all_data) + + @classmethod + def from_pandas(cls, df: pd.DataFrame) -> 'AbbreviationsCollection': + """ + Создание коллекции аббревиатур из pandas DataFrame. + """ + all_data = [] + for _, row in df.iterrows(): + try: + abbreviation = Abbreviation( + short=row['short'], + full=row['full'], + document_id=row['document_id'], + ) + all_data.append(abbreviation) + except Exception as e: + logger.warning( + f'Failed to create abbreviation from row: {row}. Error: {e}' + ) + continue + + logger.info(f'Created abbreviations collection with {len(all_data)} items') + logger.debug( + 'First 5 abbreviations: %s', ', '.join(str(abbr) for abbr in all_data[:5]) + ) + return cls(all_data) diff --git a/components/parser/docx_to_xml.py b/components/parser/docx_to_xml.py new file mode 100644 index 0000000000000000000000000000000000000000..7dd2ec34d9c91ace6046404f4da857adfa6063d0 --- /dev/null +++ b/components/parser/docx_to_xml.py @@ -0,0 +1,51 @@ +import logging +import zipfile +from pathlib import Path + +logger = logging.getLogger(__name__) + +class DocxToXml: + def __init__(self, docx_path: str): + """ + Initialize the converter with path to DOCX file + + Args: + docx_path (str): Path to the DOCX file + """ + self.docx_path = Path(docx_path) + if not self.docx_path.exists(): + raise FileNotFoundError(f"File not found: {docx_path}") + + def extract_document_xml(self) -> str: + """ + Extract document.xml content from the DOCX file + + Returns: + str: Content of document.xml file + + Raises: + ValueError: If document.xml is not found in the DOCX file + """ + try: + with zipfile.ZipFile(self.docx_path) as docx_zip: + # The main document content is always stored in word/document.xml + xml_content = docx_zip.read('word/document.xml') + return xml_content.decode('utf-8') + except KeyError: + raise ValueError("document.xml not found in the DOCX file") + except Exception as e: + raise Exception(f"Error extracting XML: {str(e)}") + + @staticmethod + def convert_file(docx_path: str) -> str: + """ + Static method to quickly convert a DOCX file to XML + + Args: + docx_path (str): Path to the DOCX file + + Returns: + str: Content of document.xml file + """ + converter = DocxToXml(docx_path) + return converter.extract_document_xml() diff --git a/components/parser/features/README.md b/components/parser/features/README.md new file mode 100644 index 0000000000000000000000000000000000000000..8289fe0b216e7c608d8e816f8dc183ec0c0cfdff --- /dev/null +++ b/components/parser/features/README.md @@ -0,0 +1,70 @@ +# Features Module + +> ВАЖНО!!! README.md сгенерировано автоматически, поэтому может содержать неточности. + +Модуль для создания структурированного датасета из обработанных документов. Включает в себя функциональность для парсинга иерархической структуры документов, создания датасета и его векторизации. + +## Основные компоненты + +### HierarchyParser + +Класс для извлечения иерархической структуры из текста документа. Позволяет: +- Парсить текстовый контент с учетом уровней вложенности +- Парсить табличный контент +- Создавать иерархическое представление документа + +### DatasetCreator + +Класс для создания структурированного датасета из обработанных документов. Функциональность: +- Обработка иерархической структуры текста и таблиц +- Создание унифицированного представления данных +- Интеграция с векторизатором для создания эмбеддингов + +### DocumentsDataset + +Класс для хранения и управления данными датасета. Возможности: +- Хранение структурированных данных документов +- Векторизация текстов с помощью предоставленного векторизатора +- Экспорт данных в pandas DataFrame +- Сохранение датасета в pickle формате + +## Структура данных + +Каждая строка датасета (`DatasetRow`) содержит следующие поля: +- Index: уникальный идентификатор строки +- Text: текстовое содержание +- DocName: имя документа +- Title: заголовок документа +- DocNumber: номер документа +- LevelParagraph: уровень параграфа +- Pargaraph: номер параграфа +- Duplicate: метка дубликата +- PartLevel1, PartLevel2: уровни частей +- Appendix: информация о приложении +- Table: информация о таблице + +## Использование + +```python +from components.embedding_extraction import EmbeddingExtractor +from components.parser.features import DatasetCreator, DocumentsDataset + +# Инициализация создателя датасета +vectorizer = EmbeddingExtractor() +creator = DatasetCreator(vectorizer) + +# Создание датасета +dataset = creator.create_dataset(parsed_xmls, hierarchies) + +# Векторизация текстов +dataset.vectorize_with(vectorizer) + +# Экспорт в pandas DataFrame +df = dataset.to_pandas() +``` + +## Зависимости + +- numpy +- pandas +- компоненты для векторизации текста (EmbeddingExtractor) \ No newline at end of file diff --git a/components/parser/features/__init__.py b/components/parser/features/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d1ed41740000229c13ffd5069d1c32efd40aea4 --- /dev/null +++ b/components/parser/features/__init__.py @@ -0,0 +1,5 @@ +from .dataset_creator import DatasetCreator +from .documents_dataset import DocumentsDataset +from .hierarchy_parser import Hierarchy, HierarchyParser + +__all__ = ['HierarchyParser', 'Hierarchy', 'DatasetCreator', 'DocumentsDataset'] diff --git a/components/parser/features/dataset_creator.py b/components/parser/features/dataset_creator.py new file mode 100644 index 0000000000000000000000000000000000000000..2130603b327d3566fcb2b41bd9e164fe0c4bc9ae --- /dev/null +++ b/components/parser/features/dataset_creator.py @@ -0,0 +1,201 @@ +import logging +import re + +from components.parser.features.documents_dataset import DatasetRow, DocumentsDataset +from components.parser.features.hierarchy_parser import Hierarchy +from components.parser.xml.structures import ParsedXML + +logger = logging.getLogger(__name__) + + +class DatasetCreator: + """ + Класс для создания датасета из обработанных документов. + """ + + def __init__( + self, + ): + """ + Инициализация создателя датасета. + """ + self._index = 0 + + def create_dataset( + self, + parsed_xmls: dict[int, ParsedXML], + hierarchies: dict[int, tuple[Hierarchy, Hierarchy]], + start_index: int = 0, + ) -> DocumentsDataset: + """ + Создание датасета из обработанных документов. + + Аргументы: + parsed_xmls: Структура с данными из XML файлов + hierarchies: Словарь с иерархическими структурами чанков + + Возвращает: + DocumentsDataset: Датасет, готовый для векторизации + """ + logger.info('Starting dataset creation from hierarchies') + self._index = start_index + dataset_rows = [] + for doc_id, (text_hierarchy, table_hierarchy) in hierarchies.items(): + xml_data = parsed_xmls[doc_id] + logger.debug( + f'Processing document {doc_id} with {len(text_hierarchy)} text sections and {len(table_hierarchy)} table sections' + ) + text_rows = self._process_text_hierarchy(text_hierarchy, xml_data) + table_rows = self._process_table_hierarchy(table_hierarchy, xml_data) + dataset_rows.extend(text_rows) + dataset_rows.extend(table_rows) + + logger.info(f'Created dataset with {len(dataset_rows)} rows') + return DocumentsDataset(dataset_rows) + + def _process_text_hierarchy( + self, + text_hierarchy: Hierarchy, + xml_data: ParsedXML, + ) -> list[DatasetRow]: + """ + Обработка иерархии текста. + """ + rows = [] + for key in text_hierarchy.keys(): + split_key = key.split('_') + + paragraph = 'unknown' + level_paragraph = 'unknown' + duplicate = 'unknown' + part_lvl1 = 'unknown' + part_lvl2 = 'unknown' + + appendix = 'unknown' + paragraph_appendix = 'unknown' + level_paragraph_appendix = 'unknown' + duplicate_appendix = 'unknown' + part_lvl1_appendix = 'unknown' + + if re.search(r'Содержание', key): + level_paragraph = -1 + paragraph = split_key[1] + elif re.search(r'Предисловие', key): + level_paragraph = -1 + paragraph = split_key[1] + if '^' in paragraph: + split_parag = paragraph.split('^') + paragraph = split_parag[0] + + # Обработка Приложений + elif re.search(r'Приложение[А-Я]\d+', key): + appendix = split_key[1].replace('Приложение', '')[0] + if len(split_key) == 3: + part_lvl1_appendix = split_key[-1] + elif len(split_key) == 4: + if 'Таблица' in key: + level_paragraph_appendix = -1 + paragraph_appendix = split_key[3] + else: + level_paragraph_appendix = split_key[2] + paragraph_appendix = split_key[3] + if ':' in paragraph_appendix: + paragraph_appendix, duplicate_appendix = ( + paragraph_appendix.split(':')[:2] + ) + paragraph_appendix = paragraph_appendix.replace( + 'PatternText', '' + ) + duplicate_appendix = duplicate_appendix.replace('Duplicate', '') + else: + paragraph_appendix = paragraph_appendix.replace( + 'PatternText', '' + ) + elif len(split_key) == 5: + level_paragraph_appendix = split_key[2] + paragraph_appendix = split_key[3] + if ':' in paragraph_appendix: + paragraph_appendix, duplicate_appendix = ( + paragraph_appendix.split(':')[:2] + ) + paragraph_appendix = paragraph_appendix.replace( + 'PatternText', '' + ) + duplicate_appendix = duplicate_appendix.replace('Duplicate', '') + else: + paragraph_appendix = paragraph_appendix.replace( + 'PatternText', '' + ) + part_lvl1_appendix = split_key[-1].replace('PartLevel', '') + else: + if len(split_key) == 2: + if '^' in split_key[1]: + split_parag = split_key[1].split('^') + level_paragraph = -1 + # paragraph = split_key[1].split('^')[-1].replace('UniqueNumber', '') + part_lvl1 = int(split_parag[1].replace('PartLevel', '')) + else: + level_paragraph = -1 + + elif len(split_key) >= 3: + level_paragraph = split_key[1][-1] + paragraph = split_key[2].replace('PatternText', '') + if ':' in paragraph: + paragraph, duplicate = paragraph.split(':')[:2] + paragraph = paragraph.replace('PatternText', '') + duplicate = duplicate.replace('Duplicate', '') + if len(split_key) == 4: + if 'Table' in key: + part_lvl1 = split_key[3] + else: + part_lvl1 = split_key[3].replace('PartLevel', '') + if len(split_key) == 5: + part_lvl1 = split_key[3].replace('PartLevel', '') + part_lvl2 = split_key[4].replace('PartLeveL', '') + + rows.append( + DatasetRow( + Index=self._index, + Text=text_hierarchy[key], + DocName=f'{xml_data.id}.XML', + DocNumber=xml_data.id, + Title=xml_data.name, + LevelParagraph=level_paragraph, + Pargaraph=paragraph, + Duplicate=duplicate, + PartLevel1=part_lvl1, + PartLevel2=part_lvl2, + Appendix=appendix, + LevelParagraphAppendix=level_paragraph_appendix, + PargaraphAppendix=paragraph_appendix, + DuplicateAppendix=duplicate_appendix, + PartLevel1Appendix=part_lvl1_appendix, + ) + ) + self._index += 1 + + return rows + + def _process_table_hierarchy( + self, + table_hierarchy: Hierarchy, + xml_data: ParsedXML, + ) -> list[DatasetRow]: + """ + Обработка иерархии таблиц. + """ + rows = [] + for key in table_hierarchy.keys(): + rows.append( + DatasetRow( + Index=self._index, + Text=table_hierarchy[key], + DocName=f'{xml_data.id}.XML', + DocNumber=xml_data.id, + Title=xml_data.name, + Table=key.split('_')[1].replace('Table', ''), + ) + ) + self._index += 1 + + return rows diff --git a/components/parser/features/documents_dataset.py b/components/parser/features/documents_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..7ee685b55d4c452521f0ce5a077418b22d63876e --- /dev/null +++ b/components/parser/features/documents_dataset.py @@ -0,0 +1,105 @@ +import logging +import pickle +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Callable + +import numpy as np +import pandas as pd + +from common.constants import UNKNOWN +from components.embedding_extraction import EmbeddingExtractor + +logger = logging.getLogger(__name__) + + +@dataclass +class DatasetRow: + """ + Класс для хранения данных одной строки датасета. + """ + + Index: int + Text: str + DocName: str + Title: str + DocNumber: str + LevelParagraph: str = field(default=UNKNOWN) + Pargaraph: str = field(default=UNKNOWN) + Duplicate: str = field(default=UNKNOWN) + PartLevel1: str = field(default=UNKNOWN) + PartLevel2: str = field(default=UNKNOWN) + Appendix: str = field(default=UNKNOWN) + LevelParagraphAppendix: str = field(default=UNKNOWN) + PargaraphAppendix: str = field(default=UNKNOWN) + DuplicateAppendix: str = field(default=UNKNOWN) + PartLevel1Appendix: str = field(default=UNKNOWN) + Table: str = field(default=UNKNOWN) + + +class DocumentsDataset: + """ + Класс для хранения данных датасета. + Содержит список строк и векторы текстов. + + Изначально не содержит векторов, чтобы запустить процесс векторизации, + нужно вызвать метод vectorize_with. + """ + + def __init__(self, rows: list[DatasetRow]): + self.rows = rows + self.vectors: np.ndarray | None = None + + def vectorize_with( + self, + vectorizer: EmbeddingExtractor, + progress_callback: Callable[[int, int], None] | None = None, + ) -> None: + """ + Векторизация текстов в датасете. + """ + logger.info('Starting dataset vectorization') + total = len(self.rows) + rows = [row.Text for row in self.rows] + vectors = vectorizer.vectorize(rows, progress_callback) + + self.vectors = vectors + logger.info(f'Completed vectorization of {total} rows') + + def to_pandas(self) -> pd.DataFrame: + """ + Преобразовать датасет в pandas DataFrame. + + Returns: + pd.DataFrame: Датафрейм с данными. + """ + df = pd.DataFrame([asdict(row) for row in self.rows]) + if self.vectors is not None: + df['Embedding'] = self.vectors.tolist() + else: + df['Embedding'] = np.nan + return df + + def to_pickle(self, path: Path) -> None: + """ + Сохранение датасета в pickle файл. + """ + logger.info(f'Saving dataset to {path}') + with open(path, 'wb') as f: + pickle.dump(self.to_pandas(), f) + logger.info('Dataset saved successfully') + + @classmethod + def from_pickle(cls, path: Path) -> 'DocumentsDataset': + """ + Загрузка датасета из pickle файла. + """ + logger.info(f'Loading dataset from {path}') + try: + with open(path, 'rb') as f: + dataset = pickle.load(f) + logger.info(f'Loaded dataset with {len(dataset.rows)} rows') + return dataset + except Exception as e: + logger.error(f'Failed to load dataset: {e}') + raise diff --git a/components/parser/features/hierarchy_parser.py b/components/parser/features/hierarchy_parser.py new file mode 100644 index 0000000000000000000000000000000000000000..0fefddcf690c8eb7527aebf8857500802b12dd75 --- /dev/null +++ b/components/parser/features/hierarchy_parser.py @@ -0,0 +1,425 @@ +from typing import List, TypeAlias +import re +import shutil +import os + + +Hierarchy: TypeAlias = dict[str, str] + + +class HierarchyParser: + def __init__(self): + self.appendix = False + self.content_flag = False + self.preface = False + self.exclude_pattern = r'^Приложение [А-Я]' + self.patterns = [ + r'^\d+\.?\s', # Соответствует "1.", "2.", и т.д. + r'^\d+\.\d+\.?\s', # Соответствует "1.1.", "2.1.", и т.д. + r'^\d+\.\d+\.\d+\.?\s', # Соответствует "1.1.1", "2.1.1", и т.д. + r'^\d+\.\d+\.\d+\.\d+\.?\s', # Соответствует "1.1.1.1", "2.1.1.1", и т.д. + r'^\d+\.\d+\.\d+\.\d+\.\d+\.?\s', # Соответствует "1.1.1.1.1", "2.1.1.1.1", и т.д. + r'^\d+\.\d+\.\d+\.\d+\.\d+\.\d+\.?\s', # Соответствует "1.1.1.1.1.1", "2.1.1.1.1.1", и т.д. + r'^\d+\.\d+\.\d+\.\d+\.\d+\.\d+\.\d+\.?\s', # Соответствует "1.1.1.1.1.1.1", и т.д. + ] + self.russian_alphabet = "АБВГДЕЁЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯабвгдеёжзийклмнопрстуфхцчшщъыьэюя" + self.english_alphabet = "ABCDEFGHIJKabcdefghijk" + self.table_name = '' + + def __init_parameters(self): + """Сбросить найденные параграфы""" + self._hierarchy = {} + self.duplicate_marker_list = [] + self.duplicate_marker_count, self.count, self.count_appendix = 0, 1, 1 + self.appendix = False + self.letter_ = None + self.content_flag = False + self.preface = False + self.table_name = '' + + def _processing_without_paragraph(self, text: str, unique_doc_number): + """ + Обрабатывает ':' если в начале нет нумерации параграфа. + Args: + text: Текст + """ + last_key = list(self._hierarchy.keys())[-1] + split_key = last_key.split('_') + if len(split_key) != 2: + last_lvl = split_key[1] + pattern_text_last = split_key[2].replace(' ', '') + key = f"{unique_doc_number}_{last_lvl}_{pattern_text_last}" + if len(split_key) == 3: + self._hierarchy[f'{key}_PartLevel{self.count}'] = f"{self._hierarchy[key]} {text}" + else: + ch = split_key[3] + self._hierarchy[f'{key}_{ch}_PartLeveL{self.count}'] = f"{self._hierarchy[key]} {text}" + else: + self._hierarchy[f'{unique_doc_number}_{text.replace("_", " ")}--DoublePoint0'] = text + + def _processing_special_markers(self, text, pattern_text, unique_doc_number): + """ + Обрабатывает маркеры 'а)', 'б)' или '-' в строке. + Args: + text: Текст + pattern_text: Паттерн для парсинга и кодирования + """ + last_key = list(self._hierarchy.keys())[-1] + split_key = last_key.split('_') + if len(split_key) != 2: + last_lvl = split_key[1] + pattern_text_last = split_key[2] + key = f'{unique_doc_number}_{last_lvl}_{pattern_text_last}' + if pattern_text is not None: + pattern_text = pattern_text.replace(')', '') + self._hierarchy[key + f'_PartLevel{pattern_text[0]}'] = f'{self._hierarchy[key]} {text}' + else: + if len(split_key) == 3: + self._hierarchy[key + f'_PartLevel{self.count}'] = f'{self._hierarchy[key]} {text}' + else: + ch = split_key[3] + if ch.replace('PartLevel', '')[-1] in self.russian_alphabet or ch.replace('PartLevel', '')[-1] in self.english_alphabet: + if len(split_key) == 4: + self.count = 1 + self._hierarchy[key + f'_{ch}_PartLeveL{self.count}'] = f'{self._hierarchy[key]} {text}' + elif re.search(r"<\d>", last_key): + self._hierarchy[key + f'_{ch}_PartLeveL{self.count}'] = f'{self._hierarchy[f"{key}_{ch}"]} {text}' + elif 'Table' in ch: + self._hierarchy[f'{key}_PartLevel{self.count}'] = f'{self._hierarchy[key]} {text}' + else: + try: + if int(ch.replace('PartLevel', '')) + 1 != self.count: + self._hierarchy[f'{key}_PartLevel{int(ch[-1]) + 1}'] = f'{self._hierarchy[key]} {text}' + else: + self._hierarchy[f'{key}_PartLevel{self.count}'] = f'{self._hierarchy[key]} {text}' + except ValueError: + self._hierarchy[f'{key}_PartLevel{self.count}'] = f'{self._hierarchy[key]} {text}' + else: + split_key = last_key.split('^') + self._hierarchy[f'{split_key[0]}^PartLevel{self.count}^UniqueNumber{self.unique_count}'] = f"{self._hierarchy[f'{split_key[0]}']} {text}" + + def _processing_appendix(self, text, level, pattern_text, unique_doc_number): + """ + Обрабатывает маркеры 'Приложение' и все, что находится внутри приложения + Args: + text: Текст + level: Уровень параграфа + pattern_text: Паттерн для парсинга и кодирования + """ + if pattern_text is not None: + pattern_text = pattern_text.replace(')', '') + if level == 13: + self.letter_ = pattern_text[-1] + self.count = 1 + self.count_appendix = 1 + if level == 13 or level is None and pattern_text is None: + if level is None: + last_key = list(self._hierarchy.keys())[-1] + split_key = last_key.split('_') + if len(split_key) == 4: + self._hierarchy[f'{last_key}_PartLevel{self.count_appendix}'] = text + self.count_appendix += 1 + elif len(split_key) == 5: + self._hierarchy[f'{"_".join(split_key[:-1])}_PartLevel{self.count_appendix}'] = text + self.count_appendix += 1 + else: + self._hierarchy[f'{unique_doc_number}_Приложение{self.letter_}{self.count}'] = text + self.count_appendix = 1 + else: + self._hierarchy[f'{unique_doc_number}_Приложение{self.letter_}{self.count}'] = text + self.count_appendix = 1 + self.count += 1 + elif level is not None and pattern_text is not None and level != 11: + key = f'{unique_doc_number}_Приложение{self.letter_}{self.count}_Level{level}_PatternText{pattern_text}' + self._hierarchy[key] = text + self.count += 1 + self.count_appendix = 1 + elif level == 10: + last_key = list(self._hierarchy.keys())[-1] + split_key = last_key.split('_') + key = f'{last_key}_PartLevel{self.count_appendix}' + if len(split_key) != 2: + self._hierarchy[key] = f"{self._hierarchy[last_key]} {text}" + else: + self._hierarchy[key] = text + self.count_appendix += 1 + elif level == 11: + last_key = list(self._hierarchy.keys())[-1] + split_key = last_key.split('_') + appendix = split_key[1] + if len(split_key) == 2: + key = f'{last_key}_PartLevel{self.count_appendix}' + self._hierarchy[key] = f"{self._hierarchy[last_key]} {text}" + elif len(split_key) != 3: + last_lvl = split_key[2] + pattern_text_last = split_key[3] + key = f'{unique_doc_number}_{appendix}_{last_lvl}_{pattern_text_last}' + self._hierarchy[f'{key}_PartLevel{self.count_appendix}'] = f"{self._hierarchy[key]} {text}" + else: + key = f'{unique_doc_number}_{appendix}' + try: + self._hierarchy[f'{key}_PartLevel{self.count_appendix}'] = f"{self._hierarchy[f'{key}_PartLevel1']} {text}" + except KeyError: + print('asdfasdf') + self._hierarchy[f'{key}_PartLevel{self.count_appendix}'] = text + self.count_appendix += 1 + + def _get_pattern(self, text): + """ + Метод находит паттерны в документе для дальнейшей обработки в соответствии с паттерном + Args: + text: Текст. + + Returns: + Код паттерна или None, Паттерн или None + + Notes: + 0-7 это разделы 1, 1.1., 1.1.1. и т.д + 10 это паттерн для поиска ':' в конце предложения + 11 это паттерн для поиска а), б) или строк начинающихся с '-', '' + 12 это паттерн для поиска наименования таблиц 'Таблица 1', 'Таблица 2' в начале строки + 13 это паттерн для поиска 'Приложение А', 'Приложение Б' в начале строки + 14 это паттерн для поиска <1>, <2> в начале строки + """ + for i, pattern in enumerate(self.patterns): + pattern = re.match(pattern, text) + if pattern: + self.preface = False + return i, pattern.group(0).replace(' ', '').replace('\xa0', '') + if re.match(r'^- ', text) or re.match(r'^– ', text): + return 11, None + if re.match(r'^ ', text): + return 11, None + if re.match(r'\d\)', text): + return 11, None + if re.match(r'\w\)', text): + pattern = re.match(r'\w\)', text) + return 11, pattern.group(0).replace('\xa0', '') + if re.match(r"<\d>", text): + pattern = re.match(r"<\d>", text) + return 14, pattern.group(0).replace('\xa0', '') + + if re.search(r':$', text) or re.search(r':$', text): + return 10, None + if re.findall(r'^Таблица \d+\.?', text): + pattern = re.match(r'^Таблица \d+\.?', text) + return 12, pattern.group(0).replace('\xa0', '') + if re.match(self.exclude_pattern, text): + pattern = re.match(self.exclude_pattern, text) + # if pattern.regs[0][1] + 2 < len(text): + # return None, None + self.appendix = True + return 13, pattern.group(0).replace('\xa0', '') + if re.match(r'^Содержание', text): + self.content_flag = True + return 15, None + if re.match(r'^Предисловие', text): + self.preface = True + return 16, None + return None, None + + def _find_duplicate_marker(self, level, pattern_text): + """ + Метод находит одинаковые параграфы в документе. И присваивает дубликату порядковый номер + Args: + level: Уровень параграфа + pattern_text: Паттерн для парсинга и кодирования + + Returns: + Название нового параграфа если дубликат был или название старого если не было дубликата + """ + if pattern_text is not None: + pattern_text = pattern_text.replace(' ', '') + if pattern_text in self.duplicate_marker_list and pattern_text[0] not in self.russian_alphabet: + if level == 0: + self.duplicate_marker_count += 1 + pattern_text = f'{pattern_text}:Duplicate{self.duplicate_marker_count}' + self.duplicate_marker_list.append(pattern_text) + return pattern_text + + def __find_last_paragraph(self, section): + for paragraph_ind in range(section.Paragraphs.Count): + paragraph = section.Paragraphs[paragraph_ind] + if paragraph.Text == '': + continue + + if paragraph.ListText != '': + text = f'{paragraph.ListText} {paragraph.Text}' + else: + text = paragraph.Text + + level_paragraph, pattern_text = self._get_pattern(text) + if pattern_text and 350 < paragraph_ind < 490 and level_paragraph < 2: + pattern_text_zxc = paragraph.Text + try: + return pattern_text_zxc + except: + return None + + def parse_table(self, doc: List, unique_doc_number): + self.__init_parameters() + flag = True + for text in doc: + text = text.strip() # удаляет пробелы в начале и конце текста параграфа. + text = text.replace('------', '').replace('--', '').replace('\u000b', ' ').replace('\t', ' ') + text = text.replace('_', ' ').replace('\u0007', ' ').replace(' ', ' ').replace('', '-') + text = text.replace(' ', ' ') + if not text: + continue + + if re.match(r'^Т\d?\d$', text) or re.match(r'^T\d?\d$', text): + try: + last_key = list(self._hierarchy.keys())[-1] + last_text = self._hierarchy[last_key] + if re.search(r': \d?\d?: $', last_text) or re.search(r': \d?\d?:$', last_text) or last_text == '': + self._hierarchy.pop(last_key, None) + except IndexError: + pass + self.table_name = text + flag = True + elif flag: + self._hierarchy[f'{unique_doc_number}_Table{self.table_name}_String{text}'] = '' + flag = False + else: + last_key = list(self._hierarchy.keys())[-1] + if text in self._hierarchy[last_key]: + continue + self._hierarchy[last_key] = f'{self._hierarchy[last_key]} {text}' + + try: + last_key = list(self._hierarchy.keys())[-1] + last_text = self._hierarchy[last_key] + if re.search(r': \d?\d?: $', last_text) or re.search(r': \d?\d?:$', last_text): + self._hierarchy.pop(last_key, None) + except IndexError: + pass + + def parse(self, doc: List, unique_doc_number, stop_appendix_list): + self.__init_parameters() + name_paragraph = None + flag = True + flag_appendix_stop = False + self.unique_count = 0 + for text in doc: + text = text.strip() # удаляет пробелы в начале и конце текста параграфа. + text = text.replace('------', '').replace('--', '').replace('\u000b', ' ').replace('\t', ' ').replace('\ufeff', '') + text = text.replace('_', ' ').replace('\u0007', ' ').replace(' ', ' ').replace('', '-').replace('', '') + if not text: + continue + + if flag: + self._hierarchy[f'{unique_doc_number}_{text.replace("_", " ")}'] = f'{text}' + flag = False + continue + + level_paragraph, pattern_text = self._get_pattern(text) # чтобы определить уровень текущего параграфа. + + if self.preface and not self.content_flag: + if level_paragraph == 16: + self._hierarchy[f'{unique_doc_number}_Предисловие'] = f'{text}' + continue + else: + self._hierarchy[f'{unique_doc_number}_Предисловие:^PatternText{pattern_text}'] = f'{self._hierarchy[f"{unique_doc_number}_Предисловие"]} {text}' + continue + + if self.content_flag: + self.preface = False + if level_paragraph == 15: + self._hierarchy[f'{unique_doc_number}_Содержание'] = f'{text}' + continue + elif level_paragraph is not None: + if level_paragraph <= 9 and '1' in pattern_text and not re.findall(r'\d+$', text): + self.content_flag = False + else: + self._hierarchy[ + f'{unique_doc_number}_Содержание'] = f'{self._hierarchy[f"{unique_doc_number}_Содержание"]} {text}' + if level_paragraph == 13: + self.appendix = False + continue + elif text not in self._hierarchy[f'{unique_doc_number}_Содержание']: + self._hierarchy[f'{unique_doc_number}_Содержание'] = f'{self._hierarchy[f"{unique_doc_number}_Содержание"]} {text}' + continue + else: + self.content_flag = False + + if self.appendix and not self.content_flag: + if level_paragraph == 13: + if pattern_text in stop_appendix_list or pattern_text[-1] in stop_appendix_list: + flag_appendix_stop = True + continue + else: + flag_appendix_stop = False + + if not flag_appendix_stop: + self._processing_appendix(text, level_paragraph, pattern_text, unique_doc_number) + continue + # TODO Проверить работоспособность и перенести в случае чего обратно перед условием про appendix + pattern_text = self._find_duplicate_marker(level_paragraph, pattern_text) + + if level_paragraph is not None: + if level_paragraph <= 9: # Пункты 1. 1.1. 2.1 и так далее до 1.1.1.1.1.1.1 + name_paragraph = text + key = f'{unique_doc_number}_Level{level_paragraph}_PatternText{pattern_text}' + self._hierarchy[key] = text + self.count = 1 + elif level_paragraph == 10: + self._processing_without_paragraph(text, unique_doc_number) + self.unique_count += 1 + self.count = 1 + elif level_paragraph == 11: + self._processing_special_markers(text, pattern_text, unique_doc_number) + self.count += 1 + elif level_paragraph == 14: + if name_paragraph is not None: + level_, pattern_text_ = self._get_pattern(name_paragraph) + self._hierarchy[f'{unique_doc_number}_Level{level_}_PatternText{pattern_text_}_PartLevel{pattern_text}'] = f'{text}' + else: + new_text = text.replace(f'{pattern_text}', '') + self._hierarchy[f'{unique_doc_number}_{new_text}'] = f'{text}' + self.count = 1 + elif level_paragraph == 12: # Обработка Таблиц + last_key = list(self._hierarchy.keys())[-self.count] + split_key = last_key.split('_') + last_lvl = split_key[1] + numer_table = pattern_text.replace('Таблица', '').replace('.', '') + if len(split_key) != 2: + pattern_text_last = last_key.split('_')[2] + key = f'{unique_doc_number}_{last_lvl}_{pattern_text_last}_Table{numer_table}' + else: + pattern_text_last = last_key.split('_')[1] + key = f'{unique_doc_number}_{pattern_text_last}_Table{numer_table}' + self._hierarchy[key] = text + self.count = 1 + else: + last_key = list(self._hierarchy.keys())[-1] + if name_paragraph is not None: + split_key = last_key.split('_') + if 'Table' in split_key[-1]: + level_paragraph, pattern_text = self._get_pattern(name_paragraph) + self._hierarchy[f'{unique_doc_number}_Level{level_paragraph}_PatternText{pattern_text}_PartLevel{self.count}'] = f'{text}' + self.count += 1 + else: + self._hierarchy[f'{last_key}'] = f'{self._hierarchy[f"{last_key}"]} {text}' + else: + self._hierarchy[f'{last_key}'] = f'{self._hierarchy[last_key]} {text}' + + def hierarchy(self): + """Вернуть иерархию документа""" + return self._hierarchy + + @staticmethod + def clear_tmp(root_path): + # Проверяем, существует ли папка + if not os.path.isdir(root_path): + raise FileNotFoundError(f"Папка {root_path} не найдена.") + + # Перебираем все файлы и папки внутри указанной директории + for item in os.listdir(root_path): + item_path = os.path.join(root_path, item) + + # Удаляем папки и их содержимое рекурсивно + if os.path.isdir(item_path): + shutil.rmtree(item_path) + else: + # Удаляем файлы + os.remove(item_path) diff --git a/components/parser/paths.py b/components/parser/paths.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea2dfabae5d2e1c0712fd8e99c66f32c98c39b --- /dev/null +++ b/components/parser/paths.py @@ -0,0 +1,27 @@ +import os +from pathlib import Path + + +class DatasetPaths: + """ + Класс для хранения путей к файлам и директориям, задействованным при формировании датасета. + + root_path: os.PathLike - корневая директория, в которой будет храниться датасет + save_intermediate_files: bool - флаг, указывающий, нужно ли сохранять промежуточные файлы, \ + которые используются только при формировании датасета + """ + + def __init__( + self, + root_path: os.PathLike, + save_intermediate_files: bool = False, + ): + self.root_path = Path(root_path) + self.abbreviations = self.root_path / 'abbreviations.csv' + self.xml_info = self.root_path / 'xml_info.csv' + self.dataset = self.root_path / 'dataset.pkl' + + self.save_intermediate_files = save_intermediate_files + self.txt_path = self.root_path / 'txt' if save_intermediate_files else None + self.txt_abbr = self.root_path / 'txt_abbr' if save_intermediate_files else None + self.jsons_path = self.root_path / 'jsons' if save_intermediate_files else None diff --git a/components/parser/pipeline.py b/components/parser/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..fa53cf6e7171f12fa14a9f18cca02927e1a87be3 --- /dev/null +++ b/components/parser/pipeline.py @@ -0,0 +1,365 @@ +import json +import logging +from collections import defaultdict +from pathlib import Path +from typing import Callable + +import numpy as np + +from components.embedding_extraction import EmbeddingExtractor +from components.parser.abbreviations.abbreviation import Abbreviation +from components.parser.abbreviations.structures import AbbreviationsCollection +from components.parser.features.dataset_creator import DatasetCreator +from components.parser.features.documents_dataset import DatasetRow, DocumentsDataset +from components.parser.features.hierarchy_parser import Hierarchy, HierarchyParser +from components.parser.paths import DatasetPaths +from components.parser.xml import ParsedXMLs, XMLParser +from components.parser.xml.constants import ACTUAL_STATUSES + +logger = logging.getLogger(__name__) + + +class DatasetCreationPipeline: + """ + Пайплайн для обработки XML файлов со следующими шагами: + 1. Парсинг XML файлов из директории + 2. Извлечение аббревиатур из распаршенного контента + 3. Применение аббревиатур к текстовому и табличному контенту + 4. Обработка контента с помощью HierarchyParser + 5. Создание и сохранение финального датасета + """ + + def __init__( + self, + dataset_id: int, + prepared_abbreviations: list[Abbreviation], + document_ids: list[int], + document_formats: list[str], + datasets_path: Path, + documents_path: Path, + vectorizer: EmbeddingExtractor | None = None, + save_intermediate_files: bool = False, + old_dataset_id: int | None = None, + ) -> None: + """ + Инициализация пайплайна. + + Args: + dataset_id: Идентификатор датасета + vectorizer: Векторизатор для создания эмбеддингов + prepared_abbreviations: Датафрейм с аббревиатурами, извлечёнными ранее + xml_ids: Список идентификаторов XML файлов + save_intermediate_files: Флаг, указывающий, нужно ли сохранять промежуточные файлы + old_dataset: Старый датасет, если он есть + """ + self.datasets_path = datasets_path + self.documents_path = documents_path + + self.dataset_id = dataset_id + self.paths = DatasetPaths( + self.datasets_path / str(dataset_id), save_intermediate_files + ) + self.document_ids = document_ids + self.document_formats = document_formats + self.prepared_abbreviations = self._group_abbreviations(prepared_abbreviations) + + self.dataset_creator = DatasetCreator() + self.vectorizer = vectorizer + self.xml_parser = XMLParser() + self.hierarchy_parser = HierarchyParser() + + self.abbreviations: AbbreviationsCollection | None = None + self.info: ParsedXMLs | None = None + self.dataset: DocumentsDataset | None = None + self.old_paths = ( + DatasetPaths( + self.datasets_path / str(old_dataset_id), + save_intermediate_files, + ) + if old_dataset_id + else None + ) + logger.info(f'DatasetCreationPipeline initialized for {dataset_id}') + + def run( + self, + progress_callback: Callable[[int, int], None] | None = None, + ) -> DocumentsDataset: + """ + Выполнение полного пайплайна обработки. + + Args: + progress_callback: Функция, которая будет вызываться при каждом шаге векторизации. + Принимает два аргумента: current и total. + current - текущий шаг векторизации. + total - общее количество шагов векторизации. + + Returns: + DocumentsDataset: Векторизованный датасет. + """ + logger.info(f'Running pipeline for {self.dataset_id}') + + # Создание выходной директории + Path(self.paths.root_path).mkdir(parents=True, exist_ok=True) + + logger.info('Folder created') + logger.info('Processing XML files...') + + # Парсинг XML файлов + parsed_xmls = self.process_xml_files() + + logger.info('XML files processed') + logger.info('Saving XML info...') + + self.info = [xml.only_info() for xml in parsed_xmls.xmls] + parsed_xmls.to_pandas().to_csv( + self.paths.xml_info, + index=False, + ) + + logger.info('XML info saved') + logger.info('Saving txt files...') + + # Сохранение промежуточных txt файлов + if self.paths.save_intermediate_files: + self._save_txt_files(parsed_xmls) + + logger.info('Txt files saved') + logger.info('Processing abbreviations...') + + # Обработка аббревиатур + self.abbreviations = self.process_abbreviations(parsed_xmls) + + logger.info('Abbreviations processed') + + logger.info('Saving abbreviations...') + AbbreviationsCollection(self.abbreviations).to_pandas().to_csv( + self.paths.abbreviations, + index=False, + ) + logger.info('Abbreviations saved') + + logger.info('Saving txt files with abbreviations...') + # Сохранение промежуточных txt файлов с применением аббревиатур + if self.paths.save_intermediate_files: + self._save_txt_files(parsed_xmls) + logger.info('Txt files with abbreviations saved') + + logger.info('Extracting hierarchies...') + hierarchies = self._extract_hierarchies(parsed_xmls) + + logger.info('Hierarchies extracted') + logger.info('Saving hierarchies...') + if self.paths.save_intermediate_files: + self._save_hierarchies(hierarchies) + logger.info('Hierarchies saved') + + logger.info('Creating dataset...') + dataset = self.create_dataset(parsed_xmls, hierarchies) + if self.vectorizer: + logger.info('Vectorizing dataset...') + dataset.vectorize_with( + self.vectorizer, + progress_callback=progress_callback, + ) + logger.info('Dataset vectorized') + + logger.info('Saving dataset...') + dataset.to_pickle(self.paths.dataset) + logger.info('Dataset saved') + return dataset + + def process_xml_files(self) -> ParsedXMLs: + """ + Парсинг XML файлов из указанной директории. + + Возвращает: + ParsedXMLs: Структура с данными из всех XML файлов + """ + parsed_xmls = [] + for document_id, document_format in zip( + self.document_ids, self.document_formats + ): + parsed_xml = XMLParser.parse( + self.documents_path / f'{document_id}.{document_format}', + include_content=True, + ) + if ('состав' in parsed_xml.name.lower()) or ( + 'составы' in parsed_xml.name.lower() + ): + continue + if parsed_xml.status not in ACTUAL_STATUSES: + continue + parsed_xml.id = document_id + parsed_xmls.append(parsed_xml) + return ParsedXMLs(parsed_xmls) + + def process_abbreviations( + self, + parsed_xmls: ParsedXMLs, + ) -> list[Abbreviation]: + """ + Обработка и применение аббревиатур к контенту документов. + + Теперь аббревиатуры уже извлечены во время парсинга, этот метод: + 1. Устанавливает document_id для извлеченных аббревиатур + 2. Применяет только документно-специфичные аббревиатуры к соответствующим документам + 3. Объединяет все аббревиатуры (извлеченные и предварительно подготовленные) для возврата + + Args: + parsed_xmls: Структура с данными из всех XML файлов + + Returns: + list[Abbreviation]: Список всех аббревиатур для датасета + """ + all_abbreviations = {} + + # Итерируем по документам + for xml in parsed_xmls.xmls: + # Устанавливаем document_id для извлеченных аббревиатур, если они есть + doc_specific_abbreviations = [] + if xml.abbreviations: + for abbreviation in xml.abbreviations: + abbreviation.document_id = xml.id + doc_specific_abbreviations = xml.abbreviations + + # Применяем только аббревиатуры, извлеченные из этого документа + if doc_specific_abbreviations: + # Если есть аббревиатуры из документа, применяем их + xml.apply_abbreviations(doc_specific_abbreviations) + + # Получаем подготовленные аббревиатуры для текущего документа + prepared_abbr = self.prepared_abbreviations.get(xml.id, []) + + # Объединяем все аббревиатуры для возврата (не для применения) + combined_abbr = (doc_specific_abbreviations or []) + prepared_abbr + + # Сохраняем объединенный список в document.abbreviations и в общем словаре + if combined_abbr: + xml.abbreviations = combined_abbr + all_abbreviations[xml.id] = combined_abbr + + return self._ungroup_abbreviations(all_abbreviations) + + def _get_already_parsed_xmls( + self, + ) -> tuple[list[int], list[DatasetRow], list[np.ndarray]]: + if self.old_paths: + self.old_dataset = DocumentsDataset.from_pickle(self.old_paths.dataset) + + ids = set([int(row.DocNumber) for row in self.old_dataset.rows]) + ids = ids.intersection(self.xml_ids) + rows = [row for row in self.old_dataset.rows if row.DocNumber in ids] + embs = [ + emb + for row, emb in zip(rows, self.old_dataset.vectors) + if row.DocNumber in ids + ] + + return ids, rows, embs + return [], [], [] + + def _extract_hierarchies( + self, + parsed_xmls: ParsedXMLs, + ) -> dict[int, tuple[Hierarchy, Hierarchy]]: + """ + Извлечение иерархических структур из текстового и табличного контента. + + Args: + parsed_xmls: Структура с данными из всех XML файлов + + Returns: + dict[int, tuple[Hierarchy, Hierarchy]]: Словарь иерархических структур для каждого документа + """ + hierarchies = {} + + for xml in parsed_xmls.xmls: + doc_id = xml.id + + # Обработка текстового контента + if xml.text: + text_lines = xml.text.to_text().split('\n') + self.hierarchy_parser.parse(text_lines, doc_id, '') + text_hierarchy = self.hierarchy_parser.hierarchy() + else: + text_hierarchy = {} + + # Обработка табличного контента + if xml.tables: + table_lines = xml.tables.to_text().split('\n') + self.hierarchy_parser.parse_table(table_lines, doc_id) + table_hierarchy = self.hierarchy_parser.hierarchy() + else: + table_hierarchy = {} + + hierarchies[doc_id] = (text_hierarchy, table_hierarchy) + + return hierarchies + + def create_dataset( + self, + parsed_xmls: ParsedXMLs, + hierarchies: dict[int, tuple[Hierarchy, Hierarchy]], + ) -> DocumentsDataset: + """ + Создание финального датасета с векторизацией. + + Args: + parsed_xmls: Структура с данными из всех XML файлов + hierarchies: Словарь с иерархической структурой документов + + Returns: + DocumentsDataset: Датасет с векторизованными текстами + """ + xmls = {xml.id: xml for xml in parsed_xmls.xmls} + self.dataset = self.dataset_creator.create_dataset(xmls, hierarchies) + return self.dataset + + def _group_abbreviations( + self, + abbreviations: list[Abbreviation], + ) -> dict[int, list[Abbreviation]]: + """ + Преобразует список аббревиатур в словарь, где ключи - идентификаторы документов, а значения - списки аббревиатур. + """ + doc_to_abbreviations = defaultdict(list) + for abbreviation in abbreviations: + doc_to_abbreviations[abbreviation.document_id].append(abbreviation) + return doc_to_abbreviations + + def _ungroup_abbreviations( + self, abbreviations: dict[int, list[Abbreviation]] + ) -> list[Abbreviation]: + """ + Преобразует словарь аббревиатур в список аббревиатур. + """ + return sum(abbreviations.values(), []) + + def _save_txt_files(self, parsed_xmls: ParsedXMLs) -> None: + """ + Сохранение текстового и табличного контента в текстовые файлы. + """ + self.paths.txt_path.mkdir(parents=True, exist_ok=True) + for xml in parsed_xmls.xmls: + with open(self.paths.txt_path / f'{xml.id}.txt', 'w', encoding='utf-8') as f: + f.write(xml.text.to_text()) + if xml.tables: + with open(self.paths.txt_path / f'{xml.id}_table.txt', 'w', encoding='utf-8') as f: + f.write(xml.tables.to_text()) + + def _save_hierarchies( + self, + hierarchies: dict[int, tuple[Hierarchy, Hierarchy]], + ) -> None: + """ + Сохранение иерархий в JSON файлы. + """ + self.paths.jsons_path.mkdir(parents=True, exist_ok=True) + for doc_id, (text_hierarchy, table_hierarchy) in hierarchies.items(): + if text_hierarchy: + with open(self.paths.jsons_path / f'{doc_id}.json', 'w', encoding='utf-8') as f: + json.dump(text_hierarchy, f) + if table_hierarchy: + with open(self.paths.jsons_path / f'{doc_id}_table.json', 'w', encoding='utf-8') as f: + json.dump(table_hierarchy, f) diff --git a/components/parser/xml/README.md b/components/parser/xml/README.md new file mode 100644 index 0000000000000000000000000000000000000000..1f0d42bc0852d57032db907bc8dbbc16a0d225ee --- /dev/null +++ b/components/parser/xml/README.md @@ -0,0 +1,85 @@ +# XML Parser + +> ВАЖНО!!! README.md сгенерировано автоматически, поэтому может содержать неточности. + +Модуль для извлечения структурированной информации из XML-файлов документов. + +## Общий процесс парсинга + +1. Рекурсивный поиск всех XML-файлов в указанной директории +2. Для каждого файла: + - Извлечение основной информации (статус, владелец, название) + - При необходимости извлечение содержимого (таблицы и текст) +3. Фильтрация документов по статусу +4. Формирование итоговой коллекции документов + +## Подробное описание компонентов + +### XMLParser + +Основной класс, координирующий работу всех парсеров. Использует три специализированных парсера: +- XMLInfoParser - для основной информации +- XMLTableParser - для таблиц +- XMLTextParser - для текстового содержимого + +### XMLInfoParser + +Извлекает основную информацию о документе: +1. Находит специальные теги (STATUS_TAG, OWNER_TAG, NAME_TAG) +2. Извлекает значения между тегами `` и `` +3. Проверяет статус документа (только USEFUL_STATUSES) +4. Формирует объект ParsedXML + +### XMLTableParser + +Обрабатывает таблицы в документе: +1. Находит все блоки таблиц (``) +2. Определяет тип таблицы (обычная/сокращения/регламентирующие документы) +3. Извлекает строки (``) и ячейки (``) +4. Обрабатывает специальные случаи (таблицы с шапкой) +5. Формирует текстовое представление таблиц + +### XMLTextParser + +Очищает текст от XML-разметки: +1. Удаляет таблицы +2. Удаляет бинарные данные +3. Удаляет специальные конструкции (KCC, CC patterns) +4. Очищает от XML-тегов +5. Обрабатывает специальные символы + +## Структуры данных + +- ParsedXML - информация об одном документе +- ParsedXMLs - коллекция документов +- ParsedTable - данные одной таблицы +- ParsedTables - коллекция таблиц +- ParsedRow - данные одной строки таблицы + +## Особенности реализации + +1. Кодировка по умолчанию: cp866 +2. Фильтрация по статусам: 'Актуальный', 'Требует актуализации', 'Упразднён' +3. Специальная обработка таблиц сокращений и регламентирующих документов +4. Возможность исключения определенных файлов при парсинге +5. Опциональное извлечение содержимого (флаг include_content) + +## Использование + +```python +from components.parser.xml.xml_parser import XMLParser + +# Создание парсера +parser = XMLParser() + +# Парсинг одного файла +result = parser.parse("path/to/file.xml", include_content=True) + +# Парсинг всех файлов в директории +results = parser.parse_all( + "path/to/dir", + encoding='cp866', + include_content=True, + ignore_files_contains=['temp', 'draft'] +) +``` \ No newline at end of file diff --git a/components/parser/xml/__init__.py b/components/parser/xml/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ac854fd9dc505997cb1676cc710cdf42caabc988 --- /dev/null +++ b/components/parser/xml/__init__.py @@ -0,0 +1,23 @@ +from .structures import ( + ParsedRow, + ParsedTable, + ParsedTables, + ParsedXML, + ParsedXMLs, +) +from .xml_info_parser import XMLInfoParser +from .xml_parser import XMLParser +from .xml_table_parser import XMLTableParser +from .xml_text_parser import XMLTextParser + +__all__ = [ + 'XMLParser', + 'XMLTableParser', + 'XMLInfoParser', + 'XMLTextParser', + 'ParsedXMLs', + 'ParsedXML', + 'ParsedTables', + 'ParsedTable', + 'ParsedRow', +] diff --git a/components/parser/xml/constants.py b/components/parser/xml/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..8d9bb6768b479c2d74a66b79363a872c16e1636a --- /dev/null +++ b/components/parser/xml/constants.py @@ -0,0 +1,39 @@ +LW_TBL: str = '' +RW_TBL: str = '' +LW_TR: str = '' +RW_TR: str = '' +LW_TC: str = '' +RW_TC: str = '' +LW_T: str = '' +RW_T: str = '' +GRID_COL: str = r'gridCol' +ABBREVIATIONS: str = 'сокращения' +ABBREVIATIONS_PATTERNS: set[str] = { + 'Сокращения', + 'Термины, определения и сокращения', + 'Обозначения и сокращения', + 'Термин', + 'Сокращение' +} +REGULATIONS: str = 'регламентирующие документы' +REGULATIONS_PATTERNS: set[str] = { + 'Инструкции были использованы следующие регламентирующие документы Компании и иные нормативные акты', + 'При разработке Положения были использованы следующие нормативные документы', + 'были использованы следующие нормативные документы', +} + +ACTUAL_STATUSES: list[str] = ['Актуальный', 'Требует актуализации'] +USEFUL_STATUSES: list[str] = ['Актуальный', 'Требует актуализации', 'Упразднён'] +STATUS_TAG: str = '$$42 Статус' +OWNER_TAGS: list[str] = [ + '$$208 БП 2-го уровня', + '$$207 БП 1-го уровня', + '$$212 Владелец документа', +] +NAME_TAG: str = '$$9 Наименование документа' +EXCLUDE_OWNERS: list[str] = [ + '$$215 Группа БП', + '$$208 БП 2-го уровня', + '$$207 БП 1-го уровня', + '$$300 GUID', +] diff --git a/components/parser/xml/structures.py b/components/parser/xml/structures.py new file mode 100644 index 0000000000000000000000000000000000000000..dd599f355e1bdd902651130cafa028e4fc878798 --- /dev/null +++ b/components/parser/xml/structures.py @@ -0,0 +1,532 @@ +import logging +from collections import Counter +from dataclasses import dataclass + +import pandas as pd + +from components.parser.abbreviations.abbreviation import Abbreviation + +logger = logging.getLogger(__name__) + + +@dataclass +class ParsedRow: + """ + Класс для хранения данных, полученных из строки таблицы. + """ + + index: int + cols: list[str] + + def apply_abbreviations(self, abbreviations: list) -> None: + """ + Применяет список аббревиатур к строке таблицы. + + Args: + abbreviations: list[Abbreviation] - список аббревиатур, которые нужно применить + """ + for abbreviation in abbreviations: + self.cols = [abbreviation.apply(column) for column in self.cols] + + def to_text(self, header: list[str] | None = None) -> str: + """ + Преобразование строки таблицы в текст. + + Пример такого преобразования: + ``` + ПиП : Привет и Пока + ``` + + Args: + header: list[str] | None - шапка таблицы, если обрабатывается многоколоночная таблица + + Returns: + str - строка таблицы в текстовом формате + """ + if header is not None: + return '\n'.join(self._apply_header(header)).strip() + else: + return ' : '.join(self.cols).strip() + + def _apply_header(self, header: list[str]) -> list[str]: + """ + Применение шапки таблицы к строке. + + Args: + header: list[str] - шапка таблицы + + Returns: + list[str] - список колонок с применённой шапкой + """ + if len(header) != len(self.cols): + logging.debug( + f'Количество колонок в строке {self.index} не совпадает с количеством колонок в шапке таблицы' + ) + named_part = [ + f'{header[col_index]}: {col_value}' + for col_index, col_value in enumerate(self.cols[: len(header)]) + ] + unnamed_part = self.cols[len(header) :] + return named_part + unnamed_part + + +@dataclass +class ParsedTable: + """ + Класс для хранения данных, полученных из таблицы. + + index: int - номер таблицы + short_type: str | None - либо "сокращения", либо "регламентирующие документы", для других таблиц не заполняется + rows: list[ParsedRow] - строки таблицы + name: str | None - название таблицы, если найдено + """ + + index: int + short_type: str | None + header: list[str] | None + rows: list[ParsedRow] + name: str | None = None + subtables: list['ParsedTable'] | None = None + note: str | None = None + rows_count: int = 0 # Количество строк в таблице + modal_cols_count: int = 0 # Модальное (самое частое) количество столбцов + has_merged_cells: bool = False # Наличие объединенных ячеек + + def apply_abbreviations(self, abbreviations) -> None: + """ + Применяет список аббревиатур ко всем элементам таблицы. + + Args: + abbreviations: list[Abbreviation] | Abbreviation - объект или список объектов аббревиатур + """ + # Преобразуем одиночную аббревиатуру в список для унификации обработки + if not isinstance(abbreviations, list): + abbreviations = [abbreviations] + + # Применяем к названию таблицы, если оно есть + if self.name: + for abbreviation in abbreviations: + self.name = abbreviation.apply(self.name) + + # Применяем к заголовку таблицы, если он есть + if self.header: + for abbreviation in abbreviations: + self.header = [abbreviation.apply(column) for column in self.header] + + # Применяем к строкам таблицы + for row in self.rows: + row.apply_abbreviations(abbreviations) + + # Применяем к примечанию, если оно есть + if self.note: + for abbreviation in abbreviations: + self.note = abbreviation.apply(self.note) + + # Применяем к подтаблицам, если они есть + if self.subtables: + for subtable in self.subtables: + subtable.apply_abbreviations(abbreviations) + + def to_text(self) -> str: + """ + Преобразование таблицы в текст для дальнейшего разбиения на чанки. + + Если таблица имеет менее 12 строк, менее 5 столбцов и не содержит объединенных ячеек, + то она будет преобразована в формат Markdown. + + Returns: + str - таблица в текстовом формате + """ + # Если таблица соответствует критериям для Markdown форматирования + if (self.rows_count < 12 and self.modal_cols_count < 5 and not self.has_merged_cells): + return self._to_markdown() + + # Иначе используем стандартный текстовый формат + result = [] + + # Основная таблица + result.append('\n\n'.join(self._rich_row(row) for row in self.rows)) + + # Подтаблицы + if self.subtables: + for subtable in self.subtables: + result.append(subtable.to_text()) + + # Примечание + if self.note: + result.append(f"Примечание к таблице {self.index + 1}: {self.note}") + + return '\n\n'.join(result) + + def _to_markdown(self) -> str: + """ + Преобразование таблицы в формат Markdown. + + Returns: + str - таблица в формате Markdown + """ + result = [] + + # Добавляем название таблицы, если оно есть + if self.name: + result.append(f"### {self.name}") + result.append("") + + # Собираем заголовок таблицы + if self.header: + header_row = "| " + " | ".join(self.header) + " |" + separator = "| " + " | ".join(["---"] * len(self.header)) + " |" + result.append(header_row) + result.append(separator) + else: + # Если нет заголовка, используем максимальное количество колонок + max_cols = max([len(row.cols) for row in self.rows]) if self.rows else 0 + if max_cols > 0: + separator = "| " + " | ".join(["---"] * max_cols) + " |" + result.append(separator) + + # Добавляем строки таблицы + for row in self.rows: + # Формируем строку в формате Markdown + markdown_row = "| " + " | ".join(row.cols) + " |" + result.append(markdown_row) + + # Добавляем примечание, если оно есть + if self.note: + result.append("") + result.append(f"*Примечание: {self.note}*") + + # Добавляем подтаблицы, если они есть + if self.subtables: + for subtable in self.subtables: + result.append("") + result.append(subtable.to_text()) + + return "\n".join(result) + + def _rich_row(self, row: ParsedRow) -> str: + """ + Преобразование строки таблицы в текст с учётом самой таблицы. + + Примеры такого преобразования: + ``` + Т1 сокращения [Название таблицы] + 1 + ПиП : Привет и Пока + ``` + ``` + Т2 [Название таблицы] + 1 + Столбец 1 : Значение 1 + Столбец 2 : Значение 2 + ``` + + Args: + row: ParsedRow - строка таблицы + + Returns: + str - строка таблицы в текстовом формате + """ + table_header = f'Т{self.index + 1}' + + if self.short_type is not None: + table_header += f' {self.short_type}' + + if self.name is not None: + table_header += f' [{self.name}]' + + return f'{table_header}\n{row.index}\n{row.to_text(self.header)}' + + def normalize(self) -> 'ParsedTable': + """ + Нормализует таблицу, обрабатывая подтаблицы и примечания. + + Нормализация включает: + 1. Определение нормального количества столбцов (мода) + 2. Выделение подтаблиц, если встречаются строки с одним столбцом, + когда нормальное число столбцов не равно 1 + 3. Обработка примечаний (последняя строка с одним столбцом) + 4. Вычисление количества строк, модального количества столбцов + 5. Определение наличия объединенных ячеек + + Returns: + ParsedTable - нормализованная таблица + """ + if not self.rows: + return self + + # Находим моду по количеству столбцов + col_counts = [len(row.cols) for row in self.rows] + mode_count = Counter(col_counts).most_common(1)[0][0] + + # Устанавливаем статистику таблицы + rows_count = len(self.rows) + modal_cols_count = mode_count + + # Проверяем наличие объединенных ячеек - если есть строки с разным количеством колонок + has_merged_cells = len(set(col_counts)) > 1 + + # Если мода не равна 1, ищем строки с одним столбцом для обработки + if mode_count != 1: + normalized_rows = [] + subtables = [] + current_subtable_rows = [] + current_subtable_name = None + + last_row_index = len(self.rows) - 1 + note = None + + for i, row in enumerate(self.rows): + if len(row.cols) == 1 and i != last_row_index: + # Это может быть подзаголовок подтаблицы + if current_subtable_rows: + # Создаем подтаблицу из накопленных строк + subtable = ParsedTable( + index=len(subtables), + short_type=self.short_type, + header=self.header, # Используем хедер основной таблицы + rows=current_subtable_rows, + name=current_subtable_name, + rows_count=len(current_subtable_rows), + modal_cols_count=mode_count, + has_merged_cells=has_merged_cells + ) + subtables.append(subtable) + + # Начинаем новую подтаблицу + # Формируем имя подтаблицы как комбинацию имени таблицы и текста подзаголовка + current_subtable_name = ( + f"{self.name}: {row.cols[0]}" if self.name else row.cols[0] + ) + current_subtable_rows = [] + elif len(row.cols) == 1 and i == last_row_index: + # Это примечание + note = row.cols[0] + else: + # Обычная строка + if current_subtable_name: + # Добавляем в текущую подтаблицу + current_subtable_rows.append(row) + else: + # Добавляем в основную таблицу + normalized_rows.append(row) + + # Добавляем последнюю подтаблицу, если она есть + if current_subtable_rows: + subtable = ParsedTable( + index=len(subtables), + short_type=self.short_type, # Используем тип основной таблицы + header=self.header, # Используем хедер основной таблицы + rows=current_subtable_rows, + name=current_subtable_name, + rows_count=len(current_subtable_rows), + modal_cols_count=mode_count, + has_merged_cells=has_merged_cells + ) + subtables.append(subtable) + + # Создаем новую таблицу с обновленными статистическими полями + return ParsedTable( + index=self.index, + short_type=self.short_type, + header=self.header, + rows=normalized_rows, + name=self.name, + subtables=subtables if subtables else None, + note=note, + rows_count=len(normalized_rows), + modal_cols_count=modal_cols_count, + has_merged_cells=has_merged_cells + ) + + # Если нет специальной обработки, просто обновляем статистические поля + self.rows_count = rows_count + self.modal_cols_count = modal_cols_count + self.has_merged_cells = has_merged_cells + return self + + +@dataclass +class ParsedTables: + """ + Класс для хранения данных, полученных из всех таблиц файла. + """ + + tables: list[ParsedTable] + + def apply_abbreviations(self, abbreviations) -> None: + """ + Применяет список аббревиатур ко всем таблицам. + + Args: + abbreviations: list[Abbreviation] | Abbreviation - объект или список объектов аббревиатур + """ + # Преобразуем одиночную аббревиатуру в список для унификации обработки + if not isinstance(abbreviations, list): + abbreviations = [abbreviations] + + for table in self.tables: + table.apply_abbreviations(abbreviations) + + def to_text(self) -> str: + """ + Преобразование всех таблиц в текст для дальнейшего разбиения на чанки. + + Returns: + str - все таблицы в текстовом формате + """ + return '\n\n'.join(table.to_text() for table in self.tables) + + def normalize(self) -> 'ParsedTables': + """ + Нормализует все таблицы, обрабатывая подтаблицы и примечания. + + Returns: + ParsedTables - нормализованные таблицы + """ + normalized_tables = [table.normalize() for table in self.tables] + return ParsedTables(tables=normalized_tables) + + +@dataclass +class ParsedText: + """ + Класс для хранения текста, полученного из XML файла. + """ + + content: list[str] + + def apply_abbreviations( + self, abbreviations: list[Abbreviation] | Abbreviation + ) -> None: + """ + Применяет список аббревиатур ко всем строкам текста. + + Args: + abbreviations: list[Abbreviation] | Abbreviation - объект или список объектов аббревиатур + """ + # Преобразуем одиночную аббревиатуру в список для унификации обработки + if not isinstance(abbreviations, list): + abbreviations = [abbreviations] + + for abbreviation in abbreviations: + self.content = [abbreviation.apply(line) for line in self.content] + + def to_text(self) -> str: + """ + Возвращает текстовое представление. + + Returns: + str - текст документа + """ + return "\n\n".join(self.content) + + +@dataclass +class ParsedXML: + """ + Класс для хранения данных, полученных из xml файла. + """ + + status: str + name: str | None + owner: str | None + filename: str + + tables: ParsedTables | None = None + text: ParsedText | None = None + abbreviations: list = None # Список аббревиатур, извлеченных из документа + + id: int | None = None + + def apply_abbreviations( + self, abbreviations: list[Abbreviation] | Abbreviation + ) -> None: + """ + Применяет список аббревиатур ко всем элементам документа. + + Args: + abbreviations: list[Abbreviation] | Abbreviation - объект или список объектов аббревиатур + """ + # Преобразуем одиночную аббревиатуру в список для унификации обработки + if not isinstance(abbreviations, list): + abbreviations = [abbreviations] + + # Применяем к содержимому таблиц, если они есть + if self.tables: + self.tables.apply_abbreviations(abbreviations) + + # Применяем к текстовому содержимому, если оно есть + if self.text: + self.text.apply_abbreviations(abbreviations) + + def apply_document_abbreviations(self) -> None: + """ + Применяет аббревиатуры, извлеченные из документа, ко всему его содержимому. + """ + if self.abbreviations: + self.apply_abbreviations(self.abbreviations) + + def __post_init__(self) -> None: + """ + Пост-инициализация объекта ParsedXML. + """ + logger.debug( + f'Initializing ParsedXML: name="{self.name}", owner="{self.owner}", status="{self.status}"' + ) + + def only_info(self) -> 'ParsedXML': + """ + Создает новый объект ParsedXML только с базовой информацией, без контента. + """ + return ParsedXML( + status=self.status, + name=self.name, + owner=self.owner, + filename=self.filename, + id=self.id, + ) + + def to_text(self) -> str: + """ + Возвращает текстовое представление всего документа, включая таблицы и текст. + + Returns: + str - полный текст документа + """ + result = [] + + # Добавляем текст таблиц, если они есть + if self.tables: + result.append(self.tables.to_text()) + + # Добавляем основной текст, если он есть + if self.text: + result.append(self.text.to_text()) + + return "\n\n".join(result) + + +@dataclass +class ParsedXMLs: + """ + Класс для хранения данных, полученных из всех xml файлов. + """ + + xmls: list[ParsedXML] + + def to_pandas(self) -> pd.DataFrame: + """ + Преобразование данных в pandas DataFrame. + """ + return pd.DataFrame( + [ + { + 'status': xml.status, + 'name': xml.name, + 'owner': xml.owner, + 'filename': xml.filename, + } + for xml in self.xmls + ] + ) diff --git a/components/parser/xml/xml_info_parser.py b/components/parser/xml/xml_info_parser.py new file mode 100644 index 0000000000000000000000000000000000000000..526229561d16f347e6197fee6544572e035bf4e7 --- /dev/null +++ b/components/parser/xml/xml_info_parser.py @@ -0,0 +1,97 @@ +import os + +from bs4 import BeautifulSoup + +from components.parser.xml.constants import ( + ACTUAL_STATUSES, + EXCLUDE_OWNERS, + NAME_TAG, + OWNER_TAGS, + STATUS_TAG, + USEFUL_STATUSES, +) +from components.parser.xml.structures import ParsedXML + + +class XMLInfoParser: + """ + Класс для парсинга основной информации из xml файлов. + """ + + def __init__(self, soup: BeautifulSoup, filepath: os.PathLike): + """ + Инициализация парсера. + + Args: + soup: BeautifulSoup - суп, содержащий xml документ + filepath: os.PathLike - путь к файлу + """ + self.filepath = filepath + self.soup = soup + + def parse(self) -> ParsedXML | None: + """ + Парсинг основной информации о xml файле. + + Returns: + ParsedXML - информация о xml файле + """ + status = self._extract_info_value(STATUS_TAG) + owner = self._extract_info_recurse(OWNER_TAGS, EXCLUDE_OWNERS) + name = self._extract_info_value(NAME_TAG) + + if (name == '') and (status == '') and (owner == ''): + status = ACTUAL_STATUSES[0] + owner = '-' + name = self.filepath.stem + + if status not in USEFUL_STATUSES: + return None + + return ParsedXML(status=status, owner=owner, name=name, filename=self.filepath) + + def _extract_info_value(self, key: str) -> str: + """ + Извлечение значения из xml по тегу с использованием BeautifulSoup. + + Args: + key: str - тег, по которому будет производиться поиск. Либо название документа, либо статус, либо владелец + + Returns: + str - значение статуса, владельца или названия документа + """ + # Ищем тег в супе + key_tag = self.soup.find(string=key) + if not key_tag: + return '' + + # Ищем следующий текстовый тег после ключевого + next_tag = key_tag.find_next('w:t') + if not next_tag: + return '' + + return next_tag.get_text() + + def _extract_info_recurse( + self, + keys: list[str], + excluding_values: list[str] | None = None, + ) -> str: + """ + Извлечение значения из xml по нескольким тегам с использованием BeautifulSoup. + + Args: + keys: list[str] - список тегов, по которым будет производиться поиск + excluding_values: list[str] | None - список значений, которые нужно исключить из результата + + Returns: + str - значение статуса, владельца или названия документа + """ + result = [] + for key in keys: + value = self._extract_info_value(key) + if excluding_values and (value in excluding_values): + continue + if value: + result.append(value) + return '/'.join(reversed(result)) diff --git a/components/parser/xml/xml_parser.py b/components/parser/xml/xml_parser.py new file mode 100644 index 0000000000000000000000000000000000000000..73ad4f1b8288e5d2e83d4d6259b5fcb35cf50e53 --- /dev/null +++ b/components/parser/xml/xml_parser.py @@ -0,0 +1,163 @@ +import logging +import os +from pathlib import Path + +from bs4 import BeautifulSoup +from tqdm import tqdm + +from components.parser.docx_to_xml import DocxToXml +from components.parser.xml.structures import ParsedXML, ParsedXMLs +from components.parser.xml.xml_info_parser import XMLInfoParser +from components.parser.xml.xml_table_parser import XMLTableParser +from components.parser.xml.xml_text_parser import XMLTextParser + +logger = logging.getLogger(__name__) + + +class XMLParser: + """ + Класс для парсинга xml файлов. + """ + + @classmethod + def parse_all( + cls, + filepath: os.PathLike, + encoding: str = 'cp866', + include_content: bool = False, + ignore_files_contains: list[str] = [], + use_tqdm: bool = False, + ) -> ParsedXMLs: + """ + Парсинг всех xml файлов в директории. + + Args: + filepath: os.PathLike - путь к директории с xml файлами + encoding: str - кодировка файлов + include_content: bool - включать ли содержимое файлов в результат + ignore_files_contains: list[str] - игнорировать файлы, содержащие эти строки в названии + use_tqdm: bool - использовать ли прогресс-бар + + Returns: + ParsedXMLs - данные, полученные из всех xml файлов + """ + files = cls._get_recursive_files(filepath, ignore_files_contains) + logger.info(f"Found {len(files)} files to parse") + if use_tqdm: + files = tqdm(files, desc='Парсинг файлов') + + parsed_xmls = [cls.parse(file, encoding, include_content) for file in files] + logger.info(f"Parsed {len(parsed_xmls)} files") + parsed_xmls = [ + xml + for xml in parsed_xmls + if ( + xml is not None + and not any( + ignore_file in xml.name for ignore_file in ignore_files_contains + ) + ) + ] + logger.info(f"Filtered {len(parsed_xmls)} files") + return ParsedXMLs(parsed_xmls) + + @classmethod + def parse( + cls, + filepath: os.PathLike, + encoding: str = 'utf-8', + include_content: bool = False, + ) -> ParsedXML | None: + """ + Парсинг xml файла. + + Args: + filepath: os.PathLike - путь к xml файлу + encoding: str - кодировка файла + include_content: bool - включать ли содержимое файла в результат + + Returns: + ParsedXML - данные, полученные из xml файла + """ + if filepath.suffix in ['.docx', '.DOCX']: + logger.info(f"Parsing docx file {filepath}") + try: + xml_text = DocxToXml(filepath).extract_document_xml() + logger.info(f"Parsed docx file {filepath}") + except Exception as e: + logger.error(f"Error parsing docx file {filepath}: {e}") + return None + else: + with open(filepath, 'r', encoding=encoding) as file: + xml_text = file.read() + + soup = BeautifulSoup(xml_text, features='xml') + + # Создаем парсер информации и получаем базовые данные + info_parser = XMLInfoParser(soup, filepath) + parsed_xml = info_parser.parse() + logger.debug(f"Parsed info for {filepath}") + + if not parsed_xml: + logger.warning(f"Failed to parse info for {filepath}") + return None + + if not include_content: + logger.debug(f"Skipping content for {filepath}") + return parsed_xml + + # Парсим таблицы и текст, сохраняя структурированные данные + table_parser = XMLTableParser(soup) + text_parser = XMLTextParser(soup) + + # Сохраняем структурированные данные вместо текста + parsed_xml.tables = table_parser.parse() + logger.debug(f"Parsed table content for {filepath}") + + parsed_xml.text = text_parser.parse() + logger.debug(f"Parsed text content for {filepath}") + + # Собираем аббревиатуры из таблиц и текста + abbreviations = [] + + # Получаем аббревиатуры из таблиц + table_abbreviations = table_parser.get_abbreviations() + if table_abbreviations: + logger.debug(f"Got {len(table_abbreviations)} abbreviations from tables") + abbreviations.extend(table_abbreviations) + + # Получаем аббревиатуры из текста + text_abbreviations = text_parser.get_abbreviations() + if text_abbreviations: + logger.debug(f"Got {len(text_abbreviations)} abbreviations from text") + abbreviations.extend(text_abbreviations) + + # Сохраняем все аббревиатуры в ParsedXML + if abbreviations: + logger.info(f"Total abbreviations extracted: {len(abbreviations)}") + parsed_xml.abbreviations = abbreviations + + # Применяем аббревиатуры к содержимому документа + parsed_xml.apply_document_abbreviations() + logger.debug(f"Applied abbreviations to document content") + + return parsed_xml + + @classmethod + def _get_recursive_files( + cls, + path_to_dir: os.PathLike, + ignore_files_contains: list[str] = [], + ) -> list[os.PathLike]: + """ + Получение всех xml файлов в директории любой вложенности. + """ + path_to_dir = Path(path_to_dir) + relative_paths = [ + path.relative_to(path_to_dir) + for path in path_to_dir.glob('**/*.xml') + if not any( + ignore_file in path.name for ignore_file in ignore_files_contains + ) + ] + return [Path(path_to_dir) / path for path in relative_paths] diff --git a/components/parser/xml/xml_table_parser.py b/components/parser/xml/xml_table_parser.py new file mode 100644 index 0000000000000000000000000000000000000000..91157e253a1eb77f8ecdc5184cc3ee0258bbd4fc --- /dev/null +++ b/components/parser/xml/xml_table_parser.py @@ -0,0 +1,262 @@ +import logging + +from bs4 import BeautifulSoup + +from components.parser.abbreviations.abbreviation import Abbreviation +from components.parser.xml.constants import (ABBREVIATIONS, + ABBREVIATIONS_PATTERNS, + REGULATIONS, REGULATIONS_PATTERNS) +from components.parser.xml.structures import (ParsedRow, ParsedTable, + ParsedTables) + +logger = logging.getLogger(__name__) + + +class XMLTableParser: + """ + Класс для парсинга таблиц из xml файлов. + """ + + def __init__(self, soup: BeautifulSoup): + self.soup = soup + self.abbreviations = [] + + def parse(self) -> ParsedTables: + """ + Парсинг таблиц из xml файла. + + Returns: + ParsedTables - все таблицы, полученные из xml файла + """ + tables = self.soup.find_all('w:tbl') + logger.info(f"Found {len(tables)} tables in XML") + + parsed_tables = [] + self.abbreviations = [] + + for table_ind, table in enumerate(tables): + table_name = self._extract_table_name(table) + + type_short = self._classify_special_types(table_name, table) + + first_row = table.find('w:tr') + columns_count = len(first_row.find_all('w:tc')) if first_row else 0 + + parsed_table = self._parse_table( + table=table, + table_index=table_ind + 1, + type_short=type_short, + use_header=columns_count != 2, + table_name=table_name, + ) + + parsed_tables.append(parsed_table) + + # Если таблица содержит сокращения, извлекаем их + if type_short == ABBREVIATIONS: + abbreviations_from_table = self._extract_abbreviations_from_table( + parsed_table + ) + if abbreviations_from_table: + self.abbreviations.extend(abbreviations_from_table) + + logger.debug(f"Parsed {len(parsed_tables)} tables") + + # Создаем и нормализуем таблицы + parsed_tables_obj = ParsedTables(tables=parsed_tables) + normalized_tables = parsed_tables_obj.normalize() + + logger.debug(f"Normalized tables: {len(normalized_tables.tables)} main tables") + + if self.abbreviations: + logger.debug( + f"Extracted {len(self.abbreviations)} abbreviations from tables" + ) + + return normalized_tables + + def get_abbreviations(self) -> list[Abbreviation]: + """ + Возвращает список аббревиатур, извлеченных из таблиц. + + Returns: + list[Abbreviation]: Список аббревиатур + """ + return self.abbreviations + + def _extract_abbreviations_from_table( + self, table: ParsedTable + ) -> list[Abbreviation]: + """ + Извлечение аббревиатур из таблицы, помеченной как "сокращения". + + Args: + table: ParsedTable - таблица сокращений + + Returns: + list[Abbreviation]: Список аббревиатур + """ + abbreviations = [] + + # Проверяем, что таблица имеет нужный формат (обычно 2 колонки) + for row in table.rows: + if len(row.cols) >= 2: + # Первая колонка обычно содержит сокращение, вторая - расшифровку + short_form = row.cols[0].strip() + full_form = row.cols[1].strip() + + # Создаем объект аббревиатуры только если оба поля не пусты + if short_form and full_form: + abbreviation = Abbreviation( + short_form=short_form, + full_form=full_form, + ) + # Обрабатываем аббревиатуру для определения типа и очистки + abbreviation.process() + abbreviations.append(abbreviation) + + return abbreviations + + @classmethod + def _parse_table( + cls, + table: BeautifulSoup, + table_index: int, + type_short: str | None, + use_header: bool = False, + table_name: str | None = None, + ) -> ParsedTable: + """ + Парсинг таблицы. + + Args: + table: BeautifulSoup - объект таблицы + table_index: int - номер таблицы в xml-файле + type_short: str | None - например, "сокращения" или "регламентирующие документы" + use_header: bool - рассматривать ли первую строку таблицы как шапку таблицы + table_name: str | None - название таблицы, если найдено + + Returns: + ParsedTable - таблица, полученная из xml файла + """ + parsed_rows = [] + header = [] if use_header else None + + rows = table.find_all('w:tr') + + for row_index, row in enumerate(rows): + columns = row.find_all('w:tc') + columns = [col.get_text() for col in columns] + + if (row_index == 0) and use_header: + header = columns + else: + parsed_rows.append(ParsedRow(index=row_index, cols=columns)) + + # Вычисляем статистические показатели таблицы + rows_count = len(parsed_rows) + + # Определяем модальное количество столбцов + if rows_count > 0: + col_counts = [len(row.cols) for row in parsed_rows] + from collections import Counter + modal_cols_count = Counter(col_counts).most_common(1)[0][0] + else: + modal_cols_count = len(header) if header else 0 + + # Инициализируем has_merged_cells как False, + # actual value will be determined in normalize method + has_merged_cells = False + + return ParsedTable( + index=table_index, + short_type=type_short, + header=header, + rows=parsed_rows, + name=table_name, + rows_count=rows_count, + modal_cols_count=modal_cols_count, + has_merged_cells=has_merged_cells, + ) + + @staticmethod + def _extract_columns_from_row( + table_row: BeautifulSoup, + ) -> list[str]: + """ + Парсинг колонок из строки таблицы. + + Args: + table_row: BeautifulSoup - объект строки таблицы + + Returns: + list[str] - список колонок, полученных из строки таблицы + """ + parsed_columns = [] + + for cell in table_row.find_all('w:tc'): + cell_text_parts = [] + + for text_element in cell.find_all('w:t'): + text_content = text_element.get_text() + + # Join all text parts from this cell and add to columns + if cell_text_parts: + parsed_columns.append(''.join(cell_text_parts)) + + return parsed_columns + + @staticmethod + def _classify_special_types( + table_name: str | None, + table: BeautifulSoup, + ) -> str | None: + """ + Поиск указаний на то, что таблица является специальной: "сокращения" или "регламентирующие документы". + + Args: + table_name: str - название таблицы + table: BeautifulSoup - объект таблицы + + Returns: + str | None - либо "сокращения", либо "регламентирующие документы", либо None, если сокращения и регламенты не найдены + """ + first_row = table.find('w:tr').text + # Проверяем наличие шаблонов в тексте перед таблицей + for pattern in ABBREVIATIONS_PATTERNS: + if (table_name and pattern.lower() in table_name.lower()) or ( + pattern in first_row.lower() + ): + return ABBREVIATIONS + + for pattern in REGULATIONS_PATTERNS: + if (table_name and pattern.lower() in table_name.lower()) or ( + pattern in first_row.lower() + ): + return REGULATIONS + + return None + + @staticmethod + def _extract_table_name( + table: BeautifulSoup, + ) -> str | None: + """ + Извлечение названия таблицы из текста перед таблицей. + + Метод ищет строки, содержащие типичные маркеры заголовков таблиц, такие как + "Таблица", "Таблица N", "Табл.", и т.д., с учетом различных вариантов написания. + + Args: + before_table_xml: str - блок xml-файла, предшествующий таблице + + Returns: + str | None - название таблицы, если найдено, иначе None + """ + # Создаем объект BeautifulSoup для парсинга XML фрагмента + previous_paragraph = table.find_previous('w:p') + + if previous_paragraph: + return previous_paragraph.get_text() + + return None diff --git a/components/parser/xml/xml_text_parser.py b/components/parser/xml/xml_text_parser.py new file mode 100644 index 0000000000000000000000000000000000000000..9dfb552f1e295d4d42d86bbea55417f4aac97684 --- /dev/null +++ b/components/parser/xml/xml_text_parser.py @@ -0,0 +1,120 @@ +import logging +import re + +from bs4 import BeautifulSoup + +from components.parser.abbreviations import AbbreviationExtractor +from components.parser.xml.structures import ParsedText + +logger = logging.getLogger(__name__) + + +class XMLTextParser: + """ + Класс для парсинга текста из xml файлов. + """ + + def __init__(self, soup: BeautifulSoup): + """ + Инициализация парсера. + + Args: + soup: BeautifulSoup - суп, содержащий весь xml документ + """ + self.soup = soup + self.abbreviation_extractor = AbbreviationExtractor() + self.abbreviations = [] + + def parse(self) -> ParsedText: + """ + Парсинг текстовой информации из xml файла. + + Returns: + ParsedText - структура с текстом, полученным из xml файла + """ + parsed_text = self._extract_text() + + # Извлекаем аббревиатуры из текста + if parsed_text and parsed_text.content: + text_content = parsed_text.to_text() + self.abbreviations = self.abbreviation_extractor.extract_abbreviations_from_text(text_content) + logger.debug(f"Extracted {len(self.abbreviations)} abbreviations from text") + + return parsed_text + + def get_abbreviations(self) -> list: + """ + Возвращает список аббревиатур, извлеченных из текста. + + Returns: + list: Список аббревиатур + """ + return self.abbreviations + + def _extract_text(self) -> ParsedText: + """ + Извлечение и очистка текста из XML. + + Returns: + ParsedText - структура, содержащая очищенный текст + """ + # Удаляем все таблицы + for table in self.soup.find_all('w:tbl'): + table.decompose() + + # Удаляем бинарные данные + for bindata in self.soup.find_all('w:bindata'): + bindata.decompose() + + # Удаляем элементы v:shape (изображения) + for shape in self.soup.find_all('v:shape'): + shape.decompose() + + # Удаляем заголовки документа + doc_props = self.soup.find('o:documentproperties') + if doc_props: + doc_props.decompose() + + # Извлекаем абзацы (теги w:p) + paragraphs = [] + for p_tag in self.soup.find_all('w:p'): + # Собираем все текстовые элементы в этом абзаце + paragraph_text_elements = [] + for text_tag in p_tag.find_all('w:t'): + content = text_tag.get_text() + + # Пропускаем специальные элементы в фигурных скобках + if content and '{' in content and '}' in content: + if '{КСС}' in content or content.startswith('{СС_'): + continue + + if content: + paragraph_text_elements.append(content) + + if paragraph_text_elements: + # Объединяем текст этого абзаца + paragraph_text = ' '.join(paragraph_text_elements) + + # Очистка текста абзаца + paragraph_text = paragraph_text.replace('&', '&') + paragraph_text = paragraph_text.replace('<', '<') + paragraph_text = paragraph_text.replace('>', '>') + paragraph_text = paragraph_text.replace('MS-Word', '') + paragraph_text = paragraph_text.replace('См. документ в ', '') + paragraph_text = paragraph_text.replace( + '------------------------------------------------------------------', '' + ) + + # Удаление фигурных скобок и их содержимого + paragraph_text = re.sub( + r'\{[\.\,\#\:\=A-Za-zа-яА-Я\d\/\s\"\-\/\?\%\_\.\&\$]+\}', '', paragraph_text + ) + + # Форматирование текста абзаца + paragraph_text = re.sub(r'[\t ]+', ' ', paragraph_text) + paragraph_text = paragraph_text.strip() + + if paragraph_text: # Добавляем только непустые абзацы + paragraphs.append(paragraph_text) + + return ParsedText(content=paragraphs) diff --git a/components/services/acronym.py b/components/services/acronym.py new file mode 100644 index 0000000000000000000000000000000000000000..90036e89ff62bf01bb26e49dd1f8bb7806890c57 --- /dev/null +++ b/components/services/acronym.py @@ -0,0 +1,154 @@ +import logging + +import pandas as pd +from sqlalchemy.orm import Session +from components.dbo.models.acronym import Acronym +from components.dbo.models.dataset import Dataset +from components.dbo.models.dataset_document import DatasetDocument +from schemas.acronym import AcronymCollectionResponse + +logger = logging.getLogger(__name__) + +class AcronymService: + """ + Сервис для работы с аббревиатурами и сокращениями. + """ + + def __init__(self, db: Session): + logger.info("Initializing AcronymService") + self.db = db + + def from_pandas(self, df: pd.DataFrame) -> None: + """ + Загрузить аббревиатуры и сокращения из pandas DataFrame. + + Args: + df: DataFrame со столбцами document_id, short_form, full_form, type + """ + logger.info(f"Loading acronyms from DataFrame with {len(df)} rows") + with self.db() as session: + try: + # Process each row in the DataFrame + for _, row in df.iterrows(): + # Create acronym + acronym = Acronym( + short_form=row['short_form'], + full_form=row['full_form'], + type=row['type'], + document_id=( + int(row['document_id']) + if pd.notna(row['document_id']) + else None + ), + ) + session.add(acronym) + + session.commit() + logger.info("Successfully loaded all acronyms") + + except Exception as e: + session.rollback() + logger.error(f"Error processing acronyms: {str(e)}") + raise e + finally: + session.close() + + def get_abbreviations(self, document_id: int) -> list[Acronym]: + """ + Получить аббревиатуры и сокращения для документа. + """ + logger.info(f"Getting abbreviations for document {document_id}") + with self.db() as session: + result = ( + session.query(Acronym) + .filter( + (Acronym.document_id == document_id) | (Acronym.document_id == None) + ) + .all() + ) + + logger.debug(f"Found {len(result)} abbreviations for document {document_id}") + return result + + def get_abbreviations_by_dataset_id(self, dataset_id: int) -> list[Acronym]: + """ + Получить аббревиатуры и сокращения для документа. + """ + logger.info(f"Getting abbreviations for dataset {dataset_id}") + return self._get_acronyms_for_dataset(dataset_id) + + def get_current_acronyms(self) -> AcronymCollectionResponse: + """ + Получить аббревиатуры и сокращения для текущего активного набора данных. + """ + logger.info("Getting acronyms for current active dataset") + with self.db() as session: + active_dataset: Dataset = session.query(Dataset).filter(Dataset.is_active == True).first() + + if not active_dataset: + logger.warning("No active dataset found") + + return AcronymCollectionResponse( + collection_id=0, + collection_name="", + collection_filename="", + updated_at=None, + acronyms={}, + ) + + result = self._get_acronyms_for_dataset(active_dataset.id) + + return AcronymCollectionResponse( + collection_id=active_dataset.id, + collection_name=active_dataset.name, + collection_filename='', + updated_at=active_dataset.date_created, #TODO: Что? + acronyms=self._compress_acronyms(result), + ) + + def _get_acronyms_for_dataset(self, dataset_id: int) -> list[Acronym]: + """ + Получить список акронимов для датасета. + + Args: + dataset_id: ID датасета + + Returns: + list[Acronym]: Список акронимов + """ + with self.db() as session: + try: + document_ids = ( + session.query(DatasetDocument.document_id) + .filter(DatasetDocument.id == dataset_id) + .all() + ) + + result = ( + session.query(Acronym) + .filter( + (Acronym.document_id.in_([doc_id[0] for doc_id in document_ids])) | (Acronym.document_id == None) + ) + .all() + ) + + logger.debug(f"Found {len(result)} acronyms for dataset {dataset_id}") + return result + finally: + pass + + def _compress_acronyms(self, acronyms: list[Acronym]) -> dict[str, list[str]]: + """ + Сжать аббревиатуры и сокращения в словарь. + """ + short_forms = {acronym.short_form for acronym in acronyms} + compressed = { + short_form: [ + acronym.full_form + for acronym in acronyms + if acronym.short_form == short_form + ] + for short_form in short_forms + } + logger.debug(f"Compressed {len(acronyms)} acronyms into {len(compressed)} unique short forms") + return compressed diff --git a/components/services/dataset.py b/components/services/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..5891b9c47c80917193e4a4f8a257d1caae736c96 --- /dev/null +++ b/components/services/dataset.py @@ -0,0 +1,680 @@ +import json +import logging +import os +import shutil +import zipfile +from datetime import datetime +from multiprocessing import Process +from pathlib import Path +from typing import Optional +from threading import Lock + +import pandas as pd +import torch +from fastapi import BackgroundTasks, HTTPException, UploadFile + +from common.common import get_source_format +from common.configuration import Configuration +from components.embedding_extraction import EmbeddingExtractor +from components.parser.features.documents_dataset import DocumentsDataset +from components.parser.pipeline import DatasetCreationPipeline +from components.parser.xml.structures import ParsedXML +from components.parser.xml.xml_parser import XMLParser +from sqlalchemy.orm import Session +from components.dbo.models.acronym import Acronym +from components.dbo.models.dataset import Dataset +from components.dbo.models.dataset_document import DatasetDocument +from components.dbo.models.document import Document +from schemas.dataset import Dataset as DatasetSchema +from schemas.dataset import DatasetExpanded as DatasetExpandedSchema +from schemas.dataset import DatasetProcessing +from schemas.dataset import DocumentsPage as DocumentsPageSchema +from schemas.dataset import SortQueryList +from schemas.document import Document as DocumentSchema +logger = logging.getLogger(__name__) + + +class DatasetService: + """ + Сервис для работы с датасетами. + """ + + def __init__( + self, + vectorizer: EmbeddingExtractor, + config: Configuration, + db: Session + ) -> None: + logger.info("DatasetService initializing") + self.db = db + self.config = config + self.parser = XMLParser() + self.vectorizer = vectorizer + self.regulations_path = Path(config.db_config.files.regulations_path) + self.documents_path = Path(config.db_config.files.documents_path) + logger.info("DatasetService initialized") + + + def get_dataset( + self, + dataset_id: int, + page: int = 1, + page_size: int = 20, + search: str = '', + sort: SortQueryList = [], + ) -> DatasetExpandedSchema: + """ + Получить пагинированную информацию о датасете и его документах. + """ + logger.info( + f"Getting dataset {dataset_id} (page={page}, size={page_size}, search='{search}')" + ) + self.raise_if_processing() + + with self.db() as session: + dataset: Dataset = ( + session.query(Dataset).filter(Dataset.id == dataset_id).first() + ) + if not dataset: + raise HTTPException(status_code=404, detail='Dataset not found') + + query = ( + session.query(Document) + .join(DatasetDocument, DatasetDocument.document_id == Document.id) + .filter(DatasetDocument.dataset_id == dataset_id) + .filter( + Document.status.in_(['Актуальный', 'Требует актуализации', 'Упразднён']) + ) + .filter(Document.title.like(f'%{search}%')) + ) + + query = self.sort_documents(query, sort) + + documents = query.offset((page - 1) * page_size).limit(page_size).all() + + total_documents = ( + session.query(Document) + .join(DatasetDocument, DatasetDocument.document_id == Document.id) + .filter(DatasetDocument.dataset_id == dataset_id) + .filter( + Document.status.in_(['Актуальный', 'Требует актуализации', 'Упразднён']) + ) + .filter(Document.title.like(f'%{search}%')) + .count() + ) + + dataset_expanded = DatasetExpandedSchema( + id=dataset.id, + name=dataset.name, + isDraft=dataset.is_draft, + isActive=dataset.is_active, + dateCreated=dataset.date_created, + data=DocumentsPageSchema( + page=[ + DocumentSchema( + id=document.id, + name=document.title, + owner=document.owner, + status=document.status, + ) + for document in documents + ], + total=total_documents, + pageNumber=page, + pageSize=page_size, + ), + ) + + return dataset_expanded + + def get_datasets(self) -> list[DatasetSchema]: + """ + Получить список всех датасетов. + """ + self.raise_if_processing() + + with self.db() as session: + datasets: list[Dataset] = session.query(Dataset).all() + return [ + DatasetSchema( + id=dataset.id, + name=dataset.name, + isDraft=dataset.is_draft, + isActive=dataset.is_active, + dateCreated=dataset.date_created + ) + for dataset in datasets + ] + + def create_draft(self, parent_id: int) -> DatasetSchema: + """ + Создать черновик датасета на основе родительского датасета. + """ + logger.info(f"Creating draft dataset from parent {parent_id}") + self.raise_if_processing() + + with self.db() as session: + parent = session.query(Dataset).filter(Dataset.id == parent_id).first() + if not parent: + raise HTTPException(status_code=404, detail='Parent dataset not found') + + if parent.is_draft: + raise HTTPException(status_code=400, detail='Parent dataset is draft') + + date = datetime.now() + dataset = Dataset( + name=f"{date.strftime('%Y-%m-%d %H:%M:%S')}", + is_draft=True, + is_active=False, + ) + + parent_documents = ( + session.query(DatasetDocument) + .filter(DatasetDocument.dataset_id == parent_id) + .all() + ) + new_dataset_documents = [ + DatasetDocument( + dataset_id=dataset.id, + document_id=document.id, + ) + for document in parent_documents + ] + + dataset.documents = new_dataset_documents + + session.add(dataset) + session.commit() + session.refresh(dataset) + + return self.get_dataset(dataset.id) + + def delete_dataset(self, dataset_id: int) -> None: + """ + Удалить черновик датасета. + """ + logger.info(f"Deleting dataset {dataset_id}") + self.raise_if_processing() + + with self.db() as session: + dataset: Dataset = session.query(Dataset).filter(Dataset.id == dataset_id).first() + + if not dataset: + raise HTTPException(status_code=404, detail='Dataset not found') + + if dataset.name == 'default': + raise HTTPException( + status_code=400, detail='Default dataset cannot be deleted' + ) + + if dataset.is_active: + raise HTTPException( + status_code=403, detail='Active dataset cannot be deleted' + ) + + session.delete(dataset) + session.commit() + + def apply_draft_task(self, dataset_id: int): + """ + Метод для выполнения в отдельном процессе. + """ + try: + with self.db() as session: + dataset = session.query(Dataset).filter(Dataset.id == dataset_id).first() + if not dataset: + raise HTTPException(status_code=404, detail=f"Dataset with id {dataset_id} not found") + + active_dataset = session.query(Dataset).filter(Dataset.is_active == True).first() + + self.apply_draft(dataset, session) + dataset.is_draft = False + dataset.is_active = True + if active_dataset: + active_dataset.is_active = False + + session.commit() + except Exception as e: + logger.error(f"Error applying draft: {e}") + raise + + + def activate_dataset(self, dataset_id: int, background_tasks: BackgroundTasks) -> DatasetExpandedSchema: + """ + Активировать датасет в фоновой задаче. + """ + + logger.info(f"Activating dataset {dataset_id}") + self.raise_if_processing() + + with self.db() as session: + dataset = ( + session.query(Dataset).filter(Dataset.id == dataset_id).first() + ) + active_dataset = session.query(Dataset).filter(Dataset.is_active).first() + if not dataset: + raise HTTPException(status_code=404, detail='Dataset not found') + + if dataset.is_active: + raise HTTPException(status_code=400, detail='Dataset is already active') + + if dataset.is_draft: + background_tasks.add_task(self.apply_draft_task, dataset_id) + else: + dataset.is_active = True + + if active_dataset: + active_dataset.is_active = False + + session.commit() + + return self.get_dataset(dataset_id) + + def get_processing(self) -> DatasetProcessing: + """ + Получить информацию о процессе обработки датасета. + """ + if Path('tmp.json').exists(): + try: + with open('tmp.json', 'r', encoding='utf-8') as f: + info = json.load(f) + except Exception as e: + logger.warning(f"Error loading processing info: {e}") + return DatasetProcessing( + status='in_progress', + total=None, + current=None, + datasetName=None, + ) + + with self.db() as session: + dataset_name = ( + session.query(Dataset) + .filter(Dataset.id == info['dataset_id']) + .first() + .name + ) + + return DatasetProcessing( + status='in_progress', + total=info['total'], + current=info['current'], + datasetName=dataset_name, + ) + + return DatasetProcessing( + status='ready', + total=None, + current=None, + datasetName=None, + ) + + def upload_zip(self, file: UploadFile) -> DatasetExpandedSchema: + """ + Загрузить архив с датасетом. + """ + logger.info(f"Uploading ZIP file {file.filename}") + self.raise_if_processing() + + file_location = Path.cwd() / 'tmp' / 'tmp.zip' + logger.debug(f"Saving uploaded file to {file_location}") + file_location.parent.mkdir(parents=True, exist_ok=True) + with open(file_location, 'wb') as f: + f.write(file.file.read()) + + with zipfile.ZipFile(file_location, 'r') as zip_ref: + zip_ref.extractall(file_location.parent) + + dataset = self.create_dataset_from_directory( + is_default=False, + directory_with_xmls=file_location.parent, + directory_with_ready_dataset=None, + ) + + file_location.unlink() + shutil.rmtree(file_location.parent) + + return self.get_dataset(dataset.id) + + def apply_draft( + self, + dataset: Dataset, + session, + ) -> None: + """ + Сохранить черновик как полноценный датасет. + """ + torch.set_num_threads(1) + logger.info(f"Applying draft dataset {dataset.id}") + if not dataset.is_draft: + logger.error(f"Dataset {dataset.id} is not a draft") + raise HTTPException( + status_code=400, detail='Dataset is not draft but trying to apply it' + ) + + TMP_PATH = Path('tmp.json') + + def progress_callback(current: int, total: int) -> None: + log_step = total // 100 + if log_step == 0: + log_step = 1 + if current % log_step != 0: + return + if (total > 10) and (current % (total // 10) == 0): + logger.info( + f"Processing dataset {dataset.id}: {current}/{total}" + ) + with open(TMP_PATH, 'w', encoding='utf-8') as f: + json.dump( + { + 'total': total, + 'current': current, + 'dataset_id': dataset.id, + }, + f, + ) + + TMP_PATH.touch() + + document_ids = [ + doc_dataset_link.document_id for doc_dataset_link in dataset.documents + ] + document_formats = [ + doc_dataset_link.document.source_format + for doc_dataset_link in dataset.documents + ] + + prepared_abbreviations = ( + session.query(Acronym).filter(Acronym.document_id.in_(document_ids)).all() + ) + + pipeline = DatasetCreationPipeline( + dataset_id=dataset.id, + vectorizer=self.vectorizer, + prepared_abbreviations=prepared_abbreviations, + document_ids=document_ids, + document_formats=document_formats, + datasets_path=self.regulations_path, + documents_path=self.documents_path, + save_intermediate_files=True, + ) + progress_callback(0, 1000) + + try: + pipeline.run(progress_callback) + except Exception as e: + logger.error(f"Error running pipeline: {e}") + raise HTTPException(status_code=500, detail=str(e)) + finally: + TMP_PATH.unlink() + + def raise_if_processing(self) -> None: + """ + Поднять ошибку, если процесс обработки датасета еще не завершен. + """ + if self.get_processing().status == 'in_progress': + logger.error("Dataset processing is already in progress") + raise HTTPException( + status_code=409, detail='Dataset processing is in progress' + ) + + def create_dataset_from_directory( + self, + is_default: bool, + directory_with_xmls: Path, + directory_with_ready_dataset: Path | None = None, + ) -> Dataset: + """ + Создать датасет из директории с xml-документами. + + Args: + is_default: Создать ли датасет по умолчанию. + directory_with_xmls: Путь к директории с xml-документами. + directory_with_processed_dataset: Путь к директории с обработанным датасетом - если не передан, будет произведена полная обработка (например, при создании датасета из скриптов). + + Returns: + Dataset: Созданный датасет. + """ + logger.info( + f"Creating {'default' if is_default else 'new'} dataset from directory {directory_with_xmls}" + ) + with self.db() as session: + documents = [] + + date = datetime.now() + name = 'default' if is_default else f'{date.strftime("%Y-%m-%d %H:%M:%S")}' + + dataset = Dataset( + name=name, + is_draft=True if directory_with_ready_dataset is None else False, + is_active=True if is_default else False, + ) + session.add(dataset) + + for subpath in self._get_recursive_dirlist(directory_with_xmls): + document, relation = self._create_document( + directory_with_xmls, subpath, dataset + ) + if document is None: + continue + documents.append(document) + session.add(document) + session.add(relation) + + logger.info(f"Created {len(documents)} documents") + + session.flush() + + if directory_with_ready_dataset is not None: + shutil.move( + directory_with_ready_dataset, + self.regulations_path / str(dataset.id), + ) + + logger.info( + f"Moved ready dataset to {self.regulations_path / str(dataset.id)}" + ) + + self.documents_path.mkdir(parents=True, exist_ok=True) + + for document in documents: + session.refresh(document) + old_filename = document.filename + new_filename = '{}.{}'.format(document.id, document.source_format) + shutil.copy( + directory_with_xmls / old_filename, self.documents_path / new_filename + ) + document.filename = new_filename + + logger.info(f"Documents renamed with ids") + + session.commit() + session.refresh(dataset) + + dataset_id = dataset.id + + + logger.info(f"Dataset {dataset_id} created") + + df = self.dataset_to_pandas(dataset_id) + + (self.regulations_path / str(dataset_id)).mkdir(parents=True, exist_ok=True) + df.to_csv( + self.regulations_path / str(dataset_id) / 'documents.csv', index=False + ) + + return dataset + + def create_empty_dataset(self, is_default: bool) -> Dataset: + """ + Создать пустой датасет. + """ + with self.db() as session: + name = ( + 'default' + if is_default + else f'{datetime.now().strftime("%Y-%m-%d %H:%M:%S")}' + ) + dataset = Dataset( + name=name, + is_active=True if is_default else False, + is_draft=False, + ) + session.add(dataset) + session.commit() + session.refresh(dataset) + + self.documents_path.mkdir(exist_ok=True) + + dataset_id = dataset.id + + + folder = self.regulations_path / str(dataset_id) + folder.mkdir(parents=True, exist_ok=True) + + pickle_creator = DocumentsDataset([]) + pickle_creator.to_pickle(folder / 'dataset.pkl') + + df = self.dataset_to_pandas(dataset_id) + df.to_csv(folder / 'documents.csv', index=False) + + return dataset + + @staticmethod + def _get_recursive_dirlist(path: Path) -> list[Path]: + """ + Возвращает список всех xml и docx файлов на всех уровнях вложенности. + + Args: + path: Путь к директории. + + Returns: + list[Path]: Список путей к xml-файлам относительно path. + """ + xml_files = set() #set для отбрасывания неуникальных путей + for ext in ('*.xml', '*.XML', '*.docx', '*.DOCX'): + xml_files.update(path.glob(f'**/{ext}')) + + return [p.relative_to(path) for p in xml_files] + + def _create_document( + self, + documents_path: Path, + subpath: os.PathLike, + dataset: Dataset, + ) -> tuple[Document | None, DatasetDocument | None]: + """ + Создаёт документ в базе данных. + + Args: + xmls_path: Путь к директории с xml-документами. + subpath: Путь к xml-документу относительно xmls_path. + dataset: Датасет, к которому относится документ. + + Returns: + tuple[Document, DatasetDocument]: Кортеж из документа и его связи с датасетом. + """ + logger.debug(f"Creating document from {subpath}") + + try: + source_format = get_source_format(str(subpath)) + parsed_xml: ParsedXML | None = self.parser.parse( + documents_path / subpath, include_content=False + ) + + if not parsed_xml: + logger.warning(f"Failed to parse file: {subpath}") + return None, None + + document = Document( + filename=str(subpath), + title=parsed_xml.name, + status=parsed_xml.status, + owner=parsed_xml.owner, + source_format=source_format, + ) + relation = DatasetDocument( + document=document, + dataset=dataset, + ) + + return document, relation + + except Exception as e: + logger.error(f"Error creating document from {subpath}: {e}") + return None, None + + def dataset_to_pandas(self, dataset_id: int) -> pd.DataFrame: + """ + Преобразовать датасет в pandas DataFrame. + """ + with self.db() as session: + links = ( + session.query(DatasetDocument) + .filter(DatasetDocument.dataset_id == dataset_id) + .all() + ) + documents = ( + session.query(Document) + .filter(Document.id.in_([link.document_id for link in links])) + .all() + ) + + return pd.DataFrame( + [ + { + 'id': document.id, + 'filename': document.filename, + 'title': document.title, + 'status': document.status, + 'owner': document.owner, + } + for document in documents + ], + columns=['id', 'filename', 'title', 'status', 'owner'], + ) + + def get_current_dataset(self) -> Dataset | None: + with self.db() as session: + print(session) + result = session.query(Dataset).filter(Dataset.is_active == True).first() + return result + + def get_default_dataset(self) -> Dataset | None: + with self.db() as session: + result = session.query(Dataset).filter(Dataset.name == 'default').first() + return result + + def sort_documents( + self, + query: "Query", # type: ignore + sort: SortQueryList, + ) -> "Query": # type: ignore + """ + Сортирует документы по заданным полям и направлениям сортировки. + """ + if sort and (len(sort.sorts) > 0): + for sort_query in sort.sorts: + field = sort_query.field + direction = sort_query.direction + + if field == 'name': + column = Document.title + elif field == 'status': + column = Document.status + elif field == 'owner': + column = Document.owner + elif field == 'id': + column = Document.id + else: + raise HTTPException( + status_code=400, detail=f'Invalid sort field: {field}' + ) + + query = query.order_by( + column.desc() if direction.lower() == 'desc' else column + ) + else: + query = query.order_by(Document.id.desc()) # Default sorting + + return query diff --git a/components/services/document.py b/components/services/document.py new file mode 100644 index 0000000000000000000000000000000000000000..b23240ab496242401c5ce657372078c39d2038c9 --- /dev/null +++ b/components/services/document.py @@ -0,0 +1,207 @@ +import logging +import os +import shutil +from pathlib import Path + +from fastapi import HTTPException, UploadFile + +from sqlalchemy.orm import Session +from common.common import get_source_format +from common.configuration import Configuration +from common.constants import PROCESSING_FORMATS +from components.parser.xml.xml_parser import XMLParser +from components.dbo.models.dataset import Dataset +from components.dbo.models.dataset_document import DatasetDocument +from components.dbo.models.document import Document +from schemas.document import Document as DocumentSchema +from schemas.document import DocumentDownload +from components.services.dataset import DatasetService + +logger = logging.getLogger(__name__) + + +class DocumentService: + """ + Сервис для работы с документами. + """ + + def __init__( + self, + dataset_service: DatasetService, + config: Configuration, + db: Session + ): + logger.info("Initializing DocumentService") + self.db = db + self.dataset_service = dataset_service + self.xml_parser = XMLParser() + self.documents_path = Path(config.db_config.files.documents_path) + + def get_document( + self, + document_id: int, + dataset_id: int | None = None, + ) -> DocumentDownload: + """ + Скачать документ по его идентификатору. + """ + logger.info(f"Getting document info for ID: {document_id}") + if dataset_id is None: + dataset_id = self.dataset_service.get_current_dataset().dataset_id + + self.dataset_service.raise_if_processing() + + with self.db() as session: + document_in_dataset = ( + session.query(DatasetDocument) + .filter( + DatasetDocument.dataset_id == dataset_id, + DatasetDocument.document_id == document_id, + ) + .first() + ) + + if not document_in_dataset: + logger.warning(f"Document not found: {document_id}") + raise HTTPException(status_code=404, detail="Document not found") + + document = ( + session.query(Document) + .filter( + Document.id == document_id, + ) + .first() + ) + + + result = DocumentDownload( + filename=f'{document.title[:40]}.{document.source_format}', + filepath=self.documents_path + / f'{document.document_id}.{document.source_format}', + ) + + logger.debug(f"Retrieved document: {result.filename}") + return result + + def add_document(self, dataset_id: int, file: UploadFile) -> DocumentSchema: + """ + Добавить документ в датасет. + """ + + self.dataset_service.raise_if_processing() + + file_location = Path.cwd() / 'tmp' / file.filename + file_location.parent.mkdir(parents=True, exist_ok=True) + with open(file_location, 'wb') as buffer: + buffer.write(file.file.read()) + + source_format = get_source_format(file.filename) + + logger.info(f"Parsing file: {file_location}") + logger.info(f"Source format: {source_format}") + + try: + parsed = self.xml_parser.parse(file_location, include_content=False) + except Exception: + raise HTTPException( + status_code=400, detail="Invalid XML file, service can't parse it" + ) + + with self.db() as session: + dataset = ( + session.query(Dataset).filter(Dataset.id == dataset_id).first() + ) + if not dataset: + raise HTTPException(status_code=404, detail='Dataset not found') + + if not dataset.is_draft: + raise HTTPException(status_code=403, detail='Dataset is not draft') + + document = Document( + title=parsed.name, + owner=parsed.owner, + status=parsed.status, + source_format=source_format, + ) + + logger.info(f"Document: {document}") + + session.add(document) + session.flush() + + logger.info(f"Document ID: {document.document_id}") + + link = DatasetDocument( + dataset_id=dataset_id, + document_id=document.document_id, + ) + session.add(link) + + if source_format in PROCESSING_FORMATS: + logger.info( + f"Moving file to: {self.documents_path / f'{document.document_id}.{source_format}'}" + ) + shutil.move( + file_location, + self.documents_path / f'{document.document_id}.{source_format}', + ) + else: + logger.error(f"Unknown source format: {source_format}") + raise HTTPException(status_code=400, detail='Unknown document format') + + if len(os.listdir(file_location.parent)) == 0: + file_location.parent.rmdir() + + session.commit() + session.refresh(document) + + result = DocumentSchema( + id=document.document_id, + name=document.title, + owner=document.owner, + status=document.status, + ) + logger.debug(f"Retrieved document: {result.name}") + return result + + def delete_document(self, dataset_id: int, document_id: int) -> None: + """ + Удалить документ из датасета. + """ + + self.dataset_service.raise_if_processing() + + with self.db() as session: + dataset_document = ( + session.query(DatasetDocument) + .filter( + DatasetDocument.dataset_id == dataset_id, + DatasetDocument.document_id == document_id, + ) + .first() + ) + + if not dataset_document: + raise HTTPException(status_code=404, detail='Document not found') + + dataset = ( + session.query(Dataset).filter(Dataset.id == dataset_id).first() + ) + + if not dataset.is_draft: + raise HTTPException(status_code=403, detail='Dataset is not draft') + + document = ( + session.query(Document).filter(Document.id == document_id).first() + ) + is_used = ( + session.query(DatasetDocument) + .filter(DatasetDocument.document_id == document_id) + .count() + ) + if is_used == 0: + os.remove(self.documents_path / f'{document_id}.{document.source_format}') + session.delete(document) + + session.delete(dataset_document) + session.commit() diff --git a/components/services/files.py b/components/services/files.py new file mode 100644 index 0000000000000000000000000000000000000000..e66b82e4dba183a4fdd6003362f72d3c047f1aee --- /dev/null +++ b/components/services/files.py @@ -0,0 +1,177 @@ +import logging +import os +import re +import subprocess +from pathlib import Path + +from docx import Document +from docx.oxml import parse_xml +from docx.oxml.ns import qn +from fastapi import HTTPException + +START_SPECIALS = re.escape(r'$$$$$SVB396') +END_SPECIALS = re.escape(r'$$18 Текст') + + +class FileService: + def __init__(self): + self.documents_path = Path(os.environ.get('DOCUMENTS_PATH', '/app/data/xmls_processed')) + self.source_path = Path(os.environ.get('SOURCE_PATH', '/app/data/SEND')) + + def prepare_file(self, filename: str) -> Path: + """ + Получает содержимое xml-файла. + Если он не обработан, файл берётся из директории SOURCE_PATH, обрабатывается и сохраняется в директорию DOCUMENTS_PATH. + + Args: + filename (str): Имя файла + + Returns: + Path: Путь к файлу + """ + file_path = self.documents_path / filename + + if not file_path.exists(): + source_file_path = self.source_path / filename + + logging.info(f"Process file: {source_file_path}") + + if (not source_file_path.exists()) or (not source_file_path.is_file()): + logging.error(f"File not found: {source_file_path}") + logging.error(f"Directory: {self.source_path} exists: {self.source_path.exists()}") + raise HTTPException(status_code=404, detail="File not found") + + with open(source_file_path, "r", encoding="utf-8") as source_file: + file_content = source_file.read() + + file_content = self._prettify_xml(file_content) + + file_path.parent.mkdir(parents=True, exist_ok=True) + + with open(file_path, "w", encoding="utf-8") as file: + file.write(file_content) + + logging.info(f"File saved: {file_path}") + + return file_path + + def prepare_pdf(self, filename: str) -> Path: + """ + Получает содержимое docx-файла. + """ + prepared_file = self.prepare_file(filename) + docx_path = prepared_file.with_suffix('.docx') + pdf_path = prepared_file.with_suffix('.pdf') + if not pdf_path.exists(): + if self._convert_to_docx(prepared_file) != 0: + raise HTTPException(status_code=400, detail="Failed to convert xml to docx") + + self._fix_style_table_docx(docx_path) + + if self._convert_to_pdf(docx_path) != 0: + raise HTTPException(status_code=400, detail="Failed to convert docx to pdf") + + docx_path.unlink() + + return pdf_path + + @staticmethod + def _prettify_xml(file_content: str) -> str: + """ + Удаляет спецсимволы из начала xml файла, чтобы документ смотрелся красиво. + + Args: + file_content (str): Содержимое xml файла + + Returns: + str: Содержимое xml файла без спецсимволов + """ + start = re.search(START_SPECIALS, file_content) + end = re.search(END_SPECIALS, file_content) + + if start and end: + return file_content[:start.start()] + file_content[end.end() + 1:] + + return file_content + + def _fix_style_table_docx(self, file_path: Path) -> None: + """ + Исправляет отображение таблиц и удаляет спецсимволы. + Args: + filename (str): Название docx файла. + """ + source_doc = Document(str(file_path)) + output_doc = Document() + + for block in source_doc.element.body: + if block.tag.endswith('p'): + clear_text = self._remove_curly_braces_content(block.text).replace('w:r>', '') + clear_text = clear_text.replace('См. документ в MS-Word', '') + output_doc.add_paragraph(clear_text) + + elif block.tag.endswith('tbl'): + old_table = parse_xml(block.xml) + old_rows = old_table.findall(qn('w:tr')) + len_old_rows = len(old_rows) + arr_lens_cell_in_rows = [len(row.findall(qn('w:tc'))) for row in old_rows] + + table = output_doc.add_table(rows=len_old_rows, cols=max(arr_lens_cell_in_rows)) + table.style = 'TableGrid' + table.autofit = False + + for ind_old_row, old_row in enumerate(old_rows): + old_cells = old_row.findall(qn('w:tc')) + for ind_old_cell, old_cell in enumerate(old_cells): + texts = old_cell.findall(f".//{qn('w:t')}") + for text in texts: + if '{' in text.text: + continue + else: + try: + table.rows[ind_old_row].cells[ind_old_cell].text = ( + table.rows[ind_old_row].cells[ind_old_cell].text + text.text + ) + except IndexError: + logging.warning('Ошибка в индексе, таблица не правильной формы') + continue + + output_doc.save(str(file_path)) + + @staticmethod + def _remove_curly_braces_content(text: str) -> str: + """Удаляет все содержимое внутри фигурных скобок, включая сами скобки.""" + return re.sub(r'\{[^{}]*\}', '', text) + + def _convert_to_pdf( + self, + file_path: Path, + ) -> int: + """ + Конвертирует docx-файл в pdf. + + Returns: + int: Код выхода. 0 - если конвертация прошла успешно. + """ + directory = str(file_path.parent) + path = str(file_path) + command = ['libreoffice', '--headless', '--convert-to', 'pdf', '--outdir', directory, path] + running = subprocess.Popen(command) + return running.wait() + + + def _convert_to_docx( + self, + file_path: Path, + ) -> int: + """ + Конвертирует xml-файл в docx. + + Returns: + int: Код выхода. 0 - если конвертация прошла успешно. + """ + directory = str(file_path.parent) + path = str(file_path) + command = ['libreoffice', '--headless', '--convert-to', 'docx', '--outdir', directory, path] + running = subprocess.Popen(command) + return running.wait() + diff --git a/components/services/llm_config.py b/components/services/llm_config.py new file mode 100644 index 0000000000000000000000000000000000000000..fa195b7160735cf3023d802d798d00a4fe3196c7 --- /dev/null +++ b/components/services/llm_config.py @@ -0,0 +1,139 @@ +import datetime +import logging + +from fastapi import HTTPException +from sqlalchemy.orm import Session + +from components.dbo.models.llm_config import LLMConfig as LLMConfigSQL +from schemas.llm_config import LLMConfig as LLMConfigScheme, LLMConfigCreateScheme + + +logger = logging.getLogger(__name__) + + +class LLMConfigService: + """ + Сервис для работы с параметрами LLM. + """ + + def __init__(self, db: Session): + logger.info("LLMConfigService initializing") + self.db = db + + + def create(self, config_scheme: LLMConfigCreateScheme): + logger.info("Creating a new config") + with self.db() as session: + new_config: LLMConfigSQL = LLMConfigSQL(**config_scheme.dict()) + session.add(new_config) + session.commit() + session.refresh(new_config) + + if(new_config.is_default): + self.set_as_default(new_config.id) + + return LLMConfigScheme(**new_config.to_dict()) + + + def get_list(self) -> list[LLMConfigScheme]: + with self.db() as session: + configs: list[LLMConfigSQL] = session.query(LLMConfigSQL).all() + + return [ + LLMConfigScheme(**config.to_dict()) + for config in configs + ] + + def get_by_id(self, id: int) -> LLMConfigScheme: + with self.db() as session: + config: LLMConfigSQL = session.query(LLMConfigSQL).filter(LLMConfigSQL.id == id).first() + + if not config: + raise HTTPException( + status_code=400, detail=f"Item with id {id} not found" + ) + + return LLMConfigScheme(**config.to_dict()) + + + def get_default(self) -> LLMConfigScheme: + with self.db() as session: + config: LLMConfigSQL = session.query(LLMConfigSQL).filter(LLMConfigSQL.is_default).first() + + if not config: + # Возвращаем дефолтнейшие параметры в случае, если ничего нет. + # Неочевидно, но в случае факапа всё работать будет. + return LLMConfigScheme( + is_default=True, + model='meta-llama/Llama-3.3-70B-Instruct-Turbo', + temperature=0.14, + top_p=0.95, + min_p=0.05, + frequency_penalty=-0.001, + presence_penalty=1.3, + n_predict=1000, + seed=42, + id=0, + date_created=datetime.datetime.now(datetime.timezone.utc) + ) + + return LLMConfigScheme(**config.to_dict()) + + + def set_as_default(self, id: int): + logger.info(f"Set default config: {id}") + + with self.db() as session: + session.query(LLMConfigSQL).filter(LLMConfigSQL.is_default).update({"is_default": False}) + config_new: LLMConfigSQL = session.query(LLMConfigSQL).filter(LLMConfigSQL.id == id).first() + + if not config_new: + raise HTTPException( + status_code=400, detail=f"Item with id {id} not found" + ) + + config_new.is_default = True + session.commit() + + + + def update(self, id: int, new_config: LLMConfigScheme): + logger.info("Updating default config") + with self.db() as session: + config: LLMConfigSQL = session.query(LLMConfigSQL).filter(LLMConfigSQL.id == id).first() + + if not config: + raise HTTPException( + status_code=400, detail=f"Item with id {id} not found" + ) + + update_data = new_config.model_dump(exclude_unset=True) + + for key, value in update_data.items(): + if hasattr(config, key): + setattr(config, key, value) + + + session.commit() + + + if(new_config.is_default): + self.set_as_default(new_config.id) + + session.refresh(config) + return config + + + def delete(self, id: int): + logger.info("Deleting config: {id}") + with self.db() as session: + config_to_del: LLMConfigSQL = session.query(LLMConfigSQL).get(id) + config_default: LLMConfigSQL = session.query(LLMConfigSQL).filter(LLMConfigSQL.is_default).first() + + if config_to_del.id == config_default.id: + raise HTTPException( + status_code=400, detail=f"The default config cannot be deleted" + ) + + session.delete(config_to_del) + session.commit() diff --git a/components/services/llm_prompt.py b/components/services/llm_prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..8dfe28c90967b21e1ff54281c4fedffa1e005bd2 --- /dev/null +++ b/components/services/llm_prompt.py @@ -0,0 +1,134 @@ +import datetime +import logging + +from fastapi import HTTPException +from sqlalchemy.orm import Session + +from components.dbo.models.llm_prompt import LlmPrompt as LlmPromptSQL +from schemas.llm_prompt import LlmPromptCreateSchema, LlmPromptSchema + + +logger = logging.getLogger(__name__) + + +class LlmPromptService: + """ + Сервис для работы с параметрами LLM. + """ + + def __init__(self, db: Session): + logger.info("LlmPromptService initializing") + self.db = db + + + def create(self, prompt_schema: LlmPromptCreateSchema): + logger.info("Creating a new prompt") + with self.db() as session: + new_prompt: LlmPromptSQL = LlmPromptSQL(**prompt_schema.model_dump()) + session.add(new_prompt) + session.commit() + session.refresh(new_prompt) + + if(new_prompt.is_default): + self.set_as_default(new_prompt.id) + + return LlmPromptSchema(**new_prompt.to_dict()) + + + def get_list(self) -> list[LlmPromptSchema]: + with self.db() as session: + prompts: list[LlmPromptSQL] = session.query(LlmPromptSQL).all() + + return [ + LlmPromptSchema(**prompt.to_dict()) + for prompt in prompts + ] + + def get_by_id(self, id: int) -> LlmPromptSchema: + with self.db() as session: + prompt: LlmPromptSQL = session.query(LlmPromptSQL).filter(LlmPromptSQL.id == id).first() + + if not prompt: + raise HTTPException( + status_code=400, detail=f"Item with id {id} not found" + ) + + return LlmPromptSchema(**prompt.to_dict()) + + + def get_default(self) -> LlmPromptSchema: + with self.db() as session: + prompt: LlmPromptSQL = session.query(LlmPromptSQL).filter(LlmPromptSQL.is_default).first() + + if not prompt: + # Возвращаем дефолтнейший промпт в случае, если ничего нет. + # Неочевидно, но в случае факапа всё работать будет. + return LlmPromptSchema( + is_default=True, + text='Ты ассистент. Ты помогаешь мне. Ты следуешь моим инструкциям.', + name='fallback', + id=0, + type="system", + date_created=datetime.datetime.now(datetime.timezone.utc) + ) + + return LlmPromptSchema(**prompt.to_dict()) + + + def set_as_default(self, id: int): + logger.info(f"Set default prompt: {id}") + + with self.db() as session: + session.query(LlmPromptSQL).filter(LlmPromptSQL.is_default).update({"is_default": False}) + prompt_new: LlmPromptSQL = session.query(LlmPromptSQL).filter(LlmPromptSQL.id == id).first() + + if not prompt_new: + raise HTTPException( + status_code=400, detail=f"Item with id {id} not found" + ) + + prompt_new.is_default = True + session.commit() + + + + def update(self, id: int, new_prompt: LlmPromptSchema): + logger.info("Updating default prompt") + with self.db() as session: + prompt: LlmPromptSQL = session.query(LlmPromptSQL).filter(LlmPromptSQL.id == id).first() + + if not prompt: + raise HTTPException( + status_code=400, detail=f"Item with id {id} not found" + ) + + update_data = new_prompt.model_dump(exclude_unset=True) + + for key, value in update_data.items(): + if hasattr(prompt, key): + setattr(prompt, key, value) + + + session.commit() + + + if(new_prompt.is_default): + self.set_as_default(new_prompt.id) + + session.refresh(prompt) + return prompt + + + def delete(self, id: int): + logger.info("Deleting prompt: {id}") + with self.db() as session: + prompt_to_del: LlmPromptSQL = session.query(LlmPromptSQL).get(id) + prompt_default: LlmPromptSQL = session.query(LlmPromptSQL).filter(LlmPromptSQL.is_default).first() + + if prompt_to_del.id == prompt_default.id: + raise HTTPException( + status_code=400, detail=f"The default prompt cannot be deleted" + ) + + session.delete(prompt_to_del) + session.commit() diff --git a/config_dev.yaml b/config_dev.yaml new file mode 100644 index 0000000000000000000000000000000000000000..aae2bcc0467cf9190c631ea1b3a9ba69ceafcbc5 --- /dev/null +++ b/config_dev.yaml @@ -0,0 +1,77 @@ +common: + log_file_path: !ENV ${LOG_FILE_PATH:./logs/common.log} + log_sql_path: !ENV ${SQLALCHEMY_DATABASE_URL:sqlite:///./logs.db} + +bd: + faiss: + model_embedding_path: !ENV ${EMBEDDING_MODEL_PATH:intfloat/multilingual-e5-large} + path_to_metadata: !ENV ${PATH_TO_METADATA:./data/regulation_datasets} + device: !ENV ${FAISS_DEVICE:cuda} + + elastic: + use_elastic: False + es_host: !ENV ${ELASTIC_HOST:localhost} + es_port: !ENV ${ELASTIC_PORT:9200} + people_path: ./data/person_card + + ranging: + use_ranging: false + alpha: 0.35 + beta: -0.15 + k_neighbors: 100 + + search: + vector_search: + use_vector_search: true + k_neighbors: 10 + + people_elastic_search: + use_people_search: false + index_name: 'people_search' + k_neighbors: 10 + + chunks_elastic_search: + use_chunks_search: true + index_name: 'nmd_full_text' + k_neighbors: 5 + + groups_elastic_search: + use_groups_search: false + index_name: 'group_search_elastic_nn' + k_neighbors: 1 + + rocks_nn_elastic_search: + use_rocks_nn_search: false + index_name: 'rocks_nn_search_elastic' + k_neighbors: 1 + + segmentation_elastic_search: + use_segmentation_search: false + index_name: 'segmentation_search_elastic' + k_neighbors: 1 + + # Если поиск будет не по чанкам, то добавить название ключа из функции search_answer словаря answer!!! + stop_index_names: ['people_answer', 'groups_answer', 'rocks_nn_answer', 'segmentation_answer'] + + abbreviation_search: + use_abbreviation_search: true + index_name: 'nmd_abbreviation_elastic' + k_neighbors: 10 + + files: + empty_start: true + regulations_path: ./data/regulation_datasets + default_regulations_path: ./data/regulation_datasets/default + documents_path: ./data/documents + +llm: + base_url: !ENV ${LLM_BASE_URL:https://api.deepinfra.com} + api_key_env: !ENV ${API_KEY_ENV:DEEPINFRA_API_KEY} + model: !ENV ${MODEL_NAME:meta-llama/Llama-3.3-70B-Instruct-Turbo} + tokenizer_name: !ENV ${TOKENIZER_NAME:unsloth/Llama-3.3-70B-Instruct} + temperature: 0.14 + top_p: 0.95 + min_p: 0.05 + frequency_penalty: -0.001 + presence_penalty: 1.3 + seed: 42 diff --git a/docker-compose-example.yaml b/docker-compose-example.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a2400fb2d9f972ef7422ab3865b67eced2e6ecc5 --- /dev/null +++ b/docker-compose-example.yaml @@ -0,0 +1,35 @@ +services: + gchat-backend: + container_name: belagro-chatbot-backend + build: + context: . + dockerfile: Dockerfile + args: + PORT: ${PORT:-8885} + environment: + - CONFIG_PATH=/app/config_dev.yaml # Конфиг + - SQLALCHEMY_DATABASE_URL=sqlite:////data/logs.db # Путь к БД + - PORT=${PORT:-8885} + - HF_HOME=/data/hf_cache + - LOG_FILE_PATH=/data/logs/common.log + - FAISS_DEVICE=cuda + - USE_ELASTIC=False + - DEEPINFRA_API_KEY=Bearer <ключ> + volumes: + - ../data:/data + - ../data/pip-cache:/root/.cache/pip + ports: + - "${PORT:-8885}:${PORT:-8885}" # Проброс порта (хост:контейнер) + networks: + - belagro-chatbot + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: 1 + capabilities: [gpu] + command: ["sh", "-c", "uvicorn main:app --host 0.0.0.0 --port ${PORT:-8885}"] +networks: + belagro-chatbot: + driver: bridge \ No newline at end of file diff --git a/main.py b/main.py new file mode 100644 index 0000000000000000000000000000000000000000..41e8b3ccc2713eecbbfc1e90140c5391986cf9f0 --- /dev/null +++ b/main.py @@ -0,0 +1,74 @@ +from contextlib import asynccontextmanager +from typing import Annotated +import dotenv +import uvicorn +import logging +import os +from fastapi import FastAPI, Depends +from fastapi.middleware.cors import CORSMiddleware + +from common.common import configure_logging +from common.configuration import Configuration +# from main_before import config + +# from routes.acronym import router as acronym_router +from common import dependencies as DI +from routes.dataset import router as dataset_router +from routes.document import router as document_router +from routes.acronym import router as acronym_router +from routes.llm import router as llm_router +from routes.llm_config import router as llm_config_router +from routes.llm_prompt import router as llm_prompt_router +from common.common import configure_logging + +# Загружаем переменные из .env +dotenv.load_dotenv() + +# from routes.feedback import router as feedback_router +# from routes.llm import router as llm_router +# from routes.log import router as log_router + +CONFIG_PATH = os.environ.get('CONFIG_PATH', 'config_dev.yaml') +print("config path: ") +print(CONFIG_PATH) +config = Configuration(CONFIG_PATH) + +logger = logging.getLogger(__name__) +configure_logging(config_file_path=config.common_config.log_file_path) + +configure_logging( + level=logging.DEBUG, + config_file_path=config.common_config.log_file_path, +) + + + +app = FastAPI(title="Assistant control panel") + +origins = ["*"] + +app.add_middleware( + CORSMiddleware, + allow_origins=origins, + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +app.include_router(llm_router) +# app.include_router(log_router) +# app.include_router(feedback_router) +app.include_router(acronym_router) +app.include_router(dataset_router) +app.include_router(document_router) +app.include_router(llm_config_router) +app.include_router(llm_prompt_router) + +if __name__ == "__main__": + uvicorn.run( + "main:app", + host="localhost", + port=8885, + reload=True, + workers=1 + ) \ No newline at end of file diff --git a/main_before_sdfsa.py b/main_before_sdfsa.py new file mode 100644 index 0000000000000000000000000000000000000000..eae2dff65a3d91479383b3b41aa244917f88b829 --- /dev/null +++ b/main_before_sdfsa.py @@ -0,0 +1,54 @@ +import logging +import os +from pathlib import Path + +import pandas as pd + +from common.common import configure_logging +from common.configuration import Configuration +from components.elastic import create_index_elastic_chunks, create_index_elastic_people +from components.embedding_extraction import EmbeddingExtractor +from controlpanel.components.datasets.dispatcher import Dispatcher +from components.nmd.services.acronym import AcronymService +from components.nmd.services.dataset import DatasetService +from components.nmd.services.document import DocumentService +from components.sqlite.create_database import create_database + +CONFIG_PATH = os.environ.get('CONFIG_PATH', './config_dev.yaml') + +config = Configuration(CONFIG_PATH) +logger = logging.getLogger(__name__) +configure_logging(config_file_path=config.common_config.log_file_path) + +logger.info(f'Start work...') +logger.info(f'Use config: {os.path.abspath(CONFIG_PATH)}') + +model = EmbeddingExtractor( + config.db_config.faiss.model_embedding_path, + config.db_config.faiss.device, +) + +dispatcher = Dispatcher(model, config, logger) + +acronym_service = AcronymService() +dataset_service = DatasetService(model, dispatcher, config) +document_service = DocumentService(dataset_service, config) + +create_database(dataset_service, config) + +current_dataset = dataset_service.get_current_dataset() + +dispatcher.reset_dataset(current_dataset.dataset_id) + +df = pd.read_pickle( + Path.cwd() + / config.db_config.files.regulations_path + / f'{current_dataset.dataset_id}' + / 'dataset.pkl' +) + +if config.db_config.elastic.use_elastic: + create_index_elastic_chunks(df, logger) + create_index_elastic_people(config.db_config.elastic.people_path, logger) + +logger.info('Loaded embedding model') diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..d0b7b1f79ec867e3452c3ee1a514c4e632e50e7d --- /dev/null +++ b/requirements.txt @@ -0,0 +1,26 @@ +urllib3==2.0 +fastapi==0.113.0 +unicorn==2.0.1.post1 +transformers==4.42.4 +pandas==2.2.2 +numpy==1.26.4 +tqdm==4.66.5 +nltk==3.8.1 +scikit-learn==1.5.1 +razdel==0.5.0 +python-dateutil==2.9.0.post0 +natasha==1.6.0 +python-docx==1.1.2 +unstructured==0.15.5 +faiss-cpu==1.8.0 +SQLAlchemy==2.0.35 +openai==1.62.0 +lxml==5.3.0 +beautifulsoup4==4.12.3 +pyaml-env==1.2.2 +torch==2.6.0 +elasticsearch==8.17.2 +uvicorn==0.34.0 +python-multipart==0.0.20 +python-dotenv==1.1.0 + diff --git a/routes/__init__.py b/routes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/routes/acronym.py b/routes/acronym.py new file mode 100644 index 0000000000000000000000000000000000000000..cd1b985f6bb7f17f698cf7eac2956cd43d68a9ef --- /dev/null +++ b/routes/acronym.py @@ -0,0 +1,19 @@ +from fastapi import APIRouter, Depends +import logging +import common.dependencies as DI +from components.services.acronym import AcronymService + +router = APIRouter() +logger = logging.getLogger(__name__) + +# Данный формат оставлен для обратной совместимости +@router.get("/collection/default") +def get_acronym_collection(acronym_service: AcronymService = Depends(DI.get_acronym_service)): + logger.info("Handling GET request to /collection/default") + try: + result = acronym_service.get_current_acronyms() + logger.info(f"Successfully retrieved acronym collection with ID {result.collection_id}") + return result + except Exception as e: + logger.error(f"Error retrieving acronym collection: {str(e)}") + raise diff --git a/routes/dataset.py b/routes/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..64d7d3c0f59cf644e60efa8c0bed7849a16146e7 --- /dev/null +++ b/routes/dataset.py @@ -0,0 +1,150 @@ +import logging +from typing import Annotated + +from fastapi import APIRouter, BackgroundTasks, HTTPException, Response, UploadFile, Depends + +from components.services.dataset import DatasetService +from schemas.dataset import (Dataset, DatasetExpanded, DatasetProcessing, + SortQuery, SortQueryList) +import common.dependencies as DI + +router = APIRouter(prefix='/datasets') +logger = logging.getLogger(__name__) + +@router.get('/') +async def get_datasets(dataset_service: Annotated[DatasetService, Depends(DI.get_dataset_service)]) -> list[Dataset]: + logger.info("Handling GET request to /datasets") + try: + result = dataset_service.get_datasets() + logger.info(f"Successfully retrieved {len(result)} datasets") + return result + except Exception as e: + logger.error(f"Error retrieving datasets: {str(e)}") + raise + + +@router.get('/processing') +async def get_processing(dataset_service: Annotated[DatasetService, Depends(DI.get_dataset_service)]) -> DatasetProcessing: + logger.info("Handling GET request to /datasets/processing") + try: + result = dataset_service.get_processing() + logger.info(f"Successfully retrieved processing status: {result.status}") + return result + except Exception as e: + logger.error(f"Error retrieving processing status: {str(e)}") + raise + + + +def try_create_default_dataset(dataset_service: DatasetService): + """ + Создаёт датасет по умолчанию, если такого нет. + """ + + if not dataset_service.get_default_dataset(): + print('creating default dataset') + if dataset_service.config.db_config.files.empty_start: + dataset_service.create_empty_dataset(is_default=True) + else: + dataset_service.create_dataset_from_directory( + is_default=True, + directory_with_xmls=dataset_service.config.db_config.files.xmls_path_default, + directory_with_ready_dataset=dataset_service.config.db_config.files.start_path, + ) + +@router.get('/try_init_default_dataset') +async def try_init_default_dataset(dataset_service: Annotated[DatasetService, Depends(DI.get_dataset_service)]): + logger.info(f"Handling GET request try_init_default_dataset") + try_create_default_dataset(dataset_service) + try: + return {"ok": True} + except Exception as e: + logger.error(f"Error creating default dataset: {str(e)}") + raise + + +@router.get('/{dataset_id}') +async def get_dataset( + dataset_id: int, + dataset_service: Annotated[DatasetService, Depends(DI.get_dataset_service)], + page: int = 1, + page_size: int = 20, + search: str = '', + sort: str = '' +) -> DatasetExpanded: + logger.info(f"Handling GET request to /datasets/{dataset_id} (page={page}, size={page_size}, search='{search}')") + + if sort: + try: + sorts = [] + for one in sort.split(','): + field, direction = one.split(':') + sorts.append(SortQuery(field=field, direction=direction)) + sort = SortQueryList(sorts=sorts) + except ValueError: + raise HTTPException( + status_code=400, + detail="Invalid sort format. Expected format: 'field:direction,field:direction'", + ) + try: + result = dataset_service.get_dataset( + dataset_id, + page=page, + page_size=page_size, + search=search, + sort=sort, + ) + logger.info(f"Successfully retrieved dataset {dataset_id}") + return result + except Exception as e: + logger.error(f"Error retrieving dataset {dataset_id}: {str(e)}") + raise e + + +@router.post('/{parent_id}/edit') +async def create_draft(parent_id: int, dataset_service: Annotated[DatasetService, Depends(DI.get_dataset_service)]) -> Dataset: + logger.info(f"Handling POST request to /datasets/{parent_id}/edit") + try: + result = dataset_service.create_draft(parent_id) + logger.info(f"Successfully created draft from dataset {parent_id}") + return result + except Exception as e: + logger.error(f"Error creating draft from dataset {parent_id}: {str(e)}") + raise e + + +@router.post('/{dataset_id}') +async def make_active(dataset_id: int, dataset_service: Annotated[DatasetService, Depends(DI.get_dataset_service)], background_tasks: BackgroundTasks) -> DatasetExpanded: + logger.info(f"Handling POST request to /datasets/{dataset_id}/activate") + try: + result = dataset_service.activate_dataset(dataset_id, background_tasks) + logger.info(f"Successfully activated dataset {dataset_id}") + return result + except Exception as e: + logger.error(f"Error activating dataset {dataset_id}: {str(e)}") + raise e + + +@router.delete('/{dataset_id}') +async def delete_dataset(dataset_id: int, dataset_service: Annotated[DatasetService, Depends(DI.get_dataset_service)]) -> None: + logger.info(f"Handling DELETE request to /datasets/{dataset_id}") + try: + dataset_service.delete_dataset(dataset_id) + logger.info(f"Successfully deleted dataset {dataset_id}") + return Response(status_code=200) + except Exception as e: + logger.error(f"Error deleting dataset {dataset_id}: {str(e)}") + raise e + + +@router.post('/') +async def upload_zip(file: UploadFile, dataset_service: Annotated[DatasetService, Depends(DI.get_dataset_service)]) -> DatasetExpanded: + logger.info(f"Handling POST request to /datasets/upload with file {file.filename}") + try: + result = dataset_service.upload_zip(file) + logger.info("Successfully uploaded and processed dataset") + return result + except Exception as e: + logger.error(f"Error uploading dataset: {str(e)}") + raise e + diff --git a/routes/document.py b/routes/document.py new file mode 100644 index 0000000000000000000000000000000000000000..30e2925197483f5a55ee6ac2c09bf763d0ff497c --- /dev/null +++ b/routes/document.py @@ -0,0 +1,42 @@ +import logging + +from fastapi import APIRouter, Response, UploadFile, Depends +from fastapi.responses import FileResponse + +from schemas.document import Document +from components.services.document import DocumentService +import common.dependencies as DI + +router = APIRouter(prefix='/datasets/{dataset_id}/documents') +logger = logging.getLogger(__name__) + + +@router.get('/{document_id}') +async def get_document(document_id: int, dataset_id: int | None = None, + document_service: DocumentService = Depends(DI.get_document_service)) -> FileResponse: + logger.info(f"Handling GET request to /documents/{document_id}") + try: + result = document_service.get_document(document_id, dataset_id) + logger.info(f"Successfully retrieved document info for ID {document_id}") + return FileResponse( + path=result.filepath, + filename=result.filename, + media_type="application/xml", + headers={"Content-Type": "application/xml; charset=cp866"}, #TODO: charset=cp866 выглядит как дичь + ) + except Exception as e: + logger.error(f"Error retrieving document {document_id}: {str(e)}") + raise + + +@router.delete('/{document_id}') +async def delete_document(dataset_id: int, document_id: int, + document_service: DocumentService = Depends(DI.get_document_service)) -> None: + document_service.delete_document(dataset_id, document_id) + return Response(status_code=200) + + +@router.post('/') +async def add_document(dataset_id: int, file: UploadFile, + document_service: DocumentService = Depends(DI.get_document_service)) -> Document: + return document_service.add_document(dataset_id, file) diff --git a/routes/feedback.py b/routes/feedback.py new file mode 100644 index 0000000000000000000000000000000000000000..9db3bc47b43f3c5851c632403f038ca116214d25 --- /dev/null +++ b/routes/feedback.py @@ -0,0 +1,78 @@ +import logging +from typing import Annotated + +from fastapi import APIRouter, Depends +from starlette import status + +from common.exceptions import FeedbackNotFoundException, LogNotFoundException +from components.dbo.models.feedback import Feedback +from components.dbo.models.log import Log +from schemas.feedback import FeedbackCreate +import common.dependencies as DI +from sqlalchemy.orm import sessionmaker + +router = APIRouter() +logger = logging.getLogger(__name__) + + +@router.get('/feedbacks', status_code=status.HTTP_200_OK) +async def get_all_feedbacks(db: Annotated[sessionmaker, Depends(DI.get_db)]): + logger.info("Handling GET request to /feedbacks") + try: + with db() as session: + result = session.query(Feedback).all() + logger.info(f"Successfully retrieved {len(result)} feedbacks") + return result + except Exception as e: + logger.error(f"Error retrieving feedbacks: {str(e)}") + raise e + + +@router.get('/feedback/{feedback_id}', status_code=status.HTTP_200_OK) +async def get_feedback(feedback_id: int, db: Annotated[sessionmaker, Depends(DI.get_db)]): + logger.info(f"Handling GET request to /feedback/{feedback_id}") + try: + with db() as session: + feedback = ( + session.query(Feedback).filter(Feedback.feedback_id == feedback_id).first() + ) + if feedback is None: + logger.warning(f"Feedback not found: {feedback_id}") + raise FeedbackNotFoundException(feedback_id) + logger.info(f"Successfully retrieved feedback {feedback_id}") + return feedback + except Exception as e: + logger.error(f"Error retrieving feedback {feedback_id}: {str(e)}") + raise e + + +@router.post('/feedback', status_code=status.HTTP_201_CREATED) +async def create_feedback(feedback: FeedbackCreate, db: Annotated[sessionmaker, Depends(DI.get_db)]): + logger.info("Handling POST request to /feedback") + try: + + with db() as session: + log_entry = session.query(Log).filter(Log.id == feedback.log_id).first() + if log_entry is None: + logger.warning(f"Log not found: {feedback.log_id}") + raise LogNotFoundException(feedback.log_id) + + new_feedback = Feedback( + log_id=feedback.log_id, + userComment=feedback.userComment, + userScore=feedback.userScore, + manualEstimate=feedback.manualEstimate, + llmEstimate=feedback.llmEstimate, + ) + + session.add(new_feedback) + session.commit() + session.refresh(new_feedback) + + logger.info( + f"Successfully created feedback with ID: {new_feedback.id}" + ) + return new_feedback + except Exception as e: + logger.error(f"Error creating feedback: {str(e)}") + raise e diff --git a/routes/file_management.py b/routes/file_management.py new file mode 100644 index 0000000000000000000000000000000000000000..750bd77b9b7c50756cbf23c2d604f1ce4e20d511 --- /dev/null +++ b/routes/file_management.py @@ -0,0 +1,42 @@ +import logging + +from fastapi import APIRouter +from fastapi.responses import FileResponse + +from components.services.files import FileService + + +router = APIRouter() + +logger = logging.getLogger(__name__) + +service = FileService() + + +@router.get("/download") +async def download_file(filename: str): + file_path = service.prepare_file(filename) + return FileResponse( + file_path, + filename=filename, + media_type="application/xml", + headers={ + "Content-Type": "application/xml; charset=cp866", + "Access-Control-Expose-Headers": "Content-Disposition" + } + ) + + +@router.get("/download_pdf") +async def download_pdf(filename: str): + file_path = service.prepare_pdf(filename) + return FileResponse( + file_path, + filename=f'{filename}.pdf', + media_type="application/pdf", + headers={ + "Content-Type": "application/pdf", + "Access-Control-Expose-Headers": "Content-Disposition" + } + ) + \ No newline at end of file diff --git a/routes/llm.py b/routes/llm.py new file mode 100644 index 0000000000000000000000000000000000000000..a8cdc51f32e2ea62bf98ebfc9040f59033f3587b --- /dev/null +++ b/routes/llm.py @@ -0,0 +1,151 @@ +import logging +from typing import Annotated, Optional, Tuple +import os +from fastapi import APIRouter, BackgroundTasks, HTTPException, Response, UploadFile, Depends +from components.llm.common import LlmParams, LlmPredictParams, Message +from components.llm.deepinfra_api import DeepInfraApi +from components.llm.llm_api import LlmApi +from components.llm.common import ChatRequest + +from common.constants import PROMPT +from components.llm.prompts import SYSTEM_PROMPT +from components.llm.utils import append_llm_response_to_history, convert_to_openai_format +from components.nmd.aggregate_answers import preprocessed_chunks +from components.nmd.llm_chunk_search import LLMChunkSearch +from components.services.dataset import DatasetService +from common.configuration import Configuration, Query, SummaryChunks +from components.datasets.dispatcher import Dispatcher +from common.exceptions import LLMResponseException +from components.dbo.models.log import Log +from components.services.llm_config import LLMConfigService +from components.services.llm_prompt import LlmPromptService +from schemas.dataset import (Dataset, DatasetExpanded, DatasetProcessing, + SortQuery, SortQueryList) +import common.dependencies as DI +from sqlalchemy.orm import Session + +router = APIRouter(prefix='/llm') +logger = logging.getLogger(__name__) + +conf = DI.get_config() +llm_params = LlmParams(**{ + "url": conf.llm_config.base_url, + "model": conf.llm_config.model, + "tokenizer": "unsloth/Llama-3.3-70B-Instruct", + "type": "deepinfra", + "default": True, + "predict_params": LlmPredictParams( + temperature=0.15, top_p=0.95, min_p=0.05, seed=42, + repetition_penalty=1.2, presence_penalty=1.1, n_predict=2000 + ), + "api_key": os.environ.get(conf.llm_config.api_key_env), + "context_length": 128000 +}) +#TODO: унести в DI +llm_api = DeepInfraApi(params=llm_params) + +@router.post("/chunks") +def get_chunks(query: Query, dispatcher: Annotated[Dispatcher, Depends(DI.get_dispatcher)]) -> SummaryChunks: + logger.info(f"Handling POST request to /chunks with query: {query.query}") + try: + result = dispatcher.search_answer(query) + logger.info("Successfully retrieved chunks") + return result + except Exception as e: + logger.error(f"Error retrieving chunks: {str(e)}") + raise e + + +def llm_answer(query: str, answer_chunks: SummaryChunks, config: Configuration + ) -> Tuple[str, str, str, int]: + """ + Метод для поиска правильного ответа с помощью LLM. + Args: + query: Запрос. + answer_chunks: Ответы векторного поиска и elastic. + + Returns: + Возвращает исходные chunks из поисков, и chunk который выбрала модель. + """ + prompt = PROMPT + llm_search = LLMChunkSearch(config.llm_config, PROMPT, logger) + return llm_search.llm_chunk_search(query, answer_chunks, prompt) + + +@router.post("/answer_llm") +def get_llm_answer(query: Query, chunks: SummaryChunks, db: Annotated[Session, Depends(DI.get_db)], config: Annotated[Configuration, Depends(DI.get_config)]): + logger.info(f"Handling POST request to /answer_llm with query: {query.query}") + try: + text_chunks, answer_llm, llm_prompt, _ = llm_answer(query.query, chunks, config) + + if not answer_llm: + logger.error("LLM returned empty response") + raise LLMResponseException() + + log_entry = Log( + llmPrompt=llm_prompt, + llmResponse=answer_llm, + userRequest=query.query, + query_type=chunks.query_type, + userName=query.userName, + ) + with db() as session: + session.add(log_entry) + session.commit() + session.refresh(log_entry) + + logger.info(f"Successfully processed LLM request, log_id: {log_entry.id}") + return { + "answer_llm": answer_llm, + "log_id": log_entry.id, + } + + except Exception as e: + logger.error(f"Error processing LLM request: {str(e)}") + raise e + + +@router.post("/chat") +async def chat(request: ChatRequest, config: Annotated[Configuration, Depends(DI.get_config)], llm_api: Annotated[DeepInfraApi, Depends(DI.get_llm_service)], prompt_service: Annotated[LlmPromptService, Depends(DI.get_llm_prompt_service)], llm_config_service: Annotated[LLMConfigService, Depends(DI.get_llm_config_service)], dispatcher: Annotated[Dispatcher, Depends(DI.get_dispatcher)]): + try: + p = llm_config_service.get_default() + system_prompt = prompt_service.get_default() + + predict_params = LlmPredictParams( + temperature=p.temperature, top_p=p.top_p, min_p=p.min_p, seed=p.seed, + frequency_penalty=p.frequency_penalty, presence_penalty=p.presence_penalty, n_predict=p.n_predict, stop=[] + ) + + #TODO: Вынести + def get_last_user_message(chat_request: ChatRequest) -> Optional[Message]: + return next( + ( + msg for msg in reversed(chat_request.history) + if msg.role == "user" and (msg.searchResults is None or not msg.searchResults) + ), + None + ) + + def insert_search_results_to_message(chat_request: ChatRequest, new_content: str) -> bool: + for msg in reversed(chat_request.history): + if msg.role == "user" and (msg.searchResults is None or not msg.searchResults): + msg.content = new_content + return True + return False + + last_query = get_last_user_message(request) + search_result = None + + if last_query: + search_result = dispatcher.search_answer(Query(query=last_query.content, query_abbreviation=last_query.content)) + text_chunks = preprocessed_chunks(search_result, None, logger) + + new_message = f'{last_query.content} /n/n{text_chunks}/n' + insert_search_results_to_message(request, new_message) + + response = await llm_api.predict_chat_stream(request, system_prompt.text, predict_params) + result = append_llm_response_to_history(request, response) + return result + except Exception as e: + logger.error(f"Error processing LLM request: {str(e)}", stack_info=True, stacklevel=10) + return {"error": str(e)} \ No newline at end of file diff --git a/routes/llm_config.py b/routes/llm_config.py new file mode 100644 index 0000000000000000000000000000000000000000..733c909ba37c647b309a0c27bc452486819133ea --- /dev/null +++ b/routes/llm_config.py @@ -0,0 +1,115 @@ +import logging + +from fastapi import APIRouter, Response, UploadFile, Depends +from fastapi.responses import FileResponse + +from schemas.llm_config import LLMConfig as LLMConfigScheme, LLMConfigCreateScheme +from components.dbo.models.llm_config import LLMConfig as LLMConfigSQL +from components.services.llm_config import LLMConfigService +import common.dependencies as DI +from schemas.llm_config import LLMConfig + +router = APIRouter(prefix='/llm_config') +logger = logging.getLogger(__name__) + + +@router.get('/') +async def get_llm_config_list(llm_config_service: LLMConfigService = Depends(DI.get_llm_config_service)) -> list[LLMConfig]: + logger.info("Handling GET request to /llm_config/{config_id}}") + try: + config = llm_config_service.get_list() + return config + except Exception as e: + logger.error(f"Error retrieving llm config: {str(e)}") + raise e + +@router.get('/default') +async def get_llm_default_config(llm_config_service: + LLMConfigService = Depends(DI.get_llm_config_service) + ) -> LLMConfig: + logger.info("Handling GET request to /llm_config/default/") + try: + config = llm_config_service.get_default() + logger.info( + f"Successfully retrieved default llm config with ID {config.id}" + ) + return config + except Exception as e: + logger.error(f"Error retrieving default llm config: {str(e)}") + raise e + + +@router.get('/{config_id}') +async def get_llm_config(config_id: int, + llm_config_service: LLMConfigService = Depends(DI.get_llm_config_service) + ) -> LLMConfig: + logger.info("Handling GET request to /llm_config/{config_id}}") + try: + config = llm_config_service.get_by_id(config_id) + logger.info( + f"Successfully retrieved llm config with ID: {config_id}" + ) + return config + except Exception as e: + logger.error(f"Error retrieving llm config: {str(e)}") + raise e + + +@router.put('/default/{config_id}') +async def set_as_default_config(config_id: int, + llm_config_service: LLMConfigService = Depends(DI.get_llm_config_service)): + logger.info("Handling PUT request to /llm_config/default/{config_id}") + try: + llm_config_service.set_as_default(config_id) + logger.info( + f"Successfully setted default llm config with ID: {config_id}" + ) + return Response(status_code=200) + except Exception as e: + logger.error(f"Error setting the default llm config: {str(e)}") + raise e + + +@router.delete('/{config_id}') +async def delete_config(config_id: int, + llm_config_service: LLMConfigService = Depends(DI.get_llm_config_service)): + logger.info("Handling DELETE request to /llm_config/{config_id}") + try: + llm_config_service.delete(config_id) + logger.info( + f"Successfully deleted llm config: {config_id}" + ) + return Response(status_code=200) + except Exception as e: + logger.error(f"Error deleting llm config: {str(e)}") + raise e + + +@router.post('/') +async def create_config(data: LLMConfigCreateScheme, + llm_config_service: LLMConfigService = Depends(DI.get_llm_config_service)): + logger.info("Handling POST request to /llm_config") + try: + new_config = llm_config_service.create(data) + logger.info( + f"Successfully created llm config with ID: {new_config.id}" + ) + return new_config + except Exception as e: + logger.error(f"Error creating llm config: {str(e)}") + raise e + + +@router.put('/{config_id}') +async def update_config(config_id: int, file: LLMConfigScheme, + llm_config_service: LLMConfigService = Depends(DI.get_llm_config_service)): + logger.info("Handling PUT request to /llm_config/{config_id}") + try: + updated_config = llm_config_service.update(config_id, file) + logger.info( + f"Successfully updated llm config with ID: {config_id}" + ) + return updated_config + except Exception as e: + logger.error(f"Error updating llm config: {str(e)}") + raise e diff --git a/routes/llm_prompt.py b/routes/llm_prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..b502f8af18caed449dc427ae67f7a6008020c18a --- /dev/null +++ b/routes/llm_prompt.py @@ -0,0 +1,113 @@ +import logging + +from fastapi import APIRouter, Response, UploadFile, Depends +from fastapi.responses import FileResponse + +from schemas.llm_prompt import LlmPromptCreateSchema, LlmPromptSchema +from components.services.llm_prompt import LlmPromptService +import common.dependencies as DI + +router = APIRouter(prefix='/llm_prompt') +logger = logging.getLogger(__name__) + + +@router.get('/') +async def get_llm_prompt_list(llm_prompt_service: LlmPromptService = Depends(DI.get_llm_prompt_service)) -> list[LlmPromptSchema]: + logger.info("Handling GET request to /llm_prompt/{prompt_id}}") + try: + prompt = llm_prompt_service.get_list() + return prompt + except Exception as e: + logger.error(f"Error retrieving llm prompt: {str(e)}") + raise e + +@router.get('/default') +async def get_llm_default_prompt(llm_prompt_service: + LlmPromptService = Depends(DI.get_llm_prompt_service) + ) -> LlmPromptSchema: + logger.info("Handling GET request to /llm_prompt/default/") + try: + prompt = llm_prompt_service.get_default() + logger.info( + f"Successfully retrieved default llm prompt with ID {prompt.id}" + ) + return prompt + except Exception as e: + logger.error(f"Error retrieving default llm prompt: {str(e)}") + raise e + + +@router.get('/{prompt_id}') +async def get_llm_prompt(prompt_id: int, + llm_prompt_service: LlmPromptService = Depends(DI.get_llm_prompt_service) + ) -> LlmPromptSchema: + logger.info("Handling GET request to /llm_prompt/{prompt_id}}") + try: + prompt = llm_prompt_service.get_by_id(prompt_id) + logger.info( + f"Successfully retrieved llm prompt with ID: {prompt_id}" + ) + return prompt + except Exception as e: + logger.error(f"Error retrieving llm prompt: {str(e)}") + raise e + + +@router.put('/default/{prompt_id}') +async def set_as_default_prompt(prompt_id: int, + llm_prompt_service: LlmPromptService = Depends(DI.get_llm_prompt_service)): + logger.info("Handling PUT request to /llm_prompt/default/{prompt_id}") + try: + llm_prompt_service.set_as_default(prompt_id) + logger.info( + f"Successfully setted default llm prompt with ID: {prompt_id}" + ) + return Response(status_code=200) + except Exception as e: + logger.error(f"Error setting the default llm prompt: {str(e)}") + raise e + + +@router.delete('/{prompt_id}') +async def delete_prompt(prompt_id: int, + llm_prompt_service: LlmPromptService = Depends(DI.get_llm_prompt_service)): + logger.info("Handling DELETE request to /llm_prompt/{prompt_id}") + try: + llm_prompt_service.delete(prompt_id) + logger.info( + f"Successfully deleted llm prompt: {prompt_id}" + ) + return Response(status_code=200) + except Exception as e: + logger.error(f"Error deleting llm prompt: {str(e)}") + raise e + + +@router.post('/') +async def create_prompt(data: LlmPromptCreateSchema, + llm_prompt_service: LlmPromptService = Depends(DI.get_llm_prompt_service)): + logger.info("Handling POST request to /llm_prompt") + try: + new_prompt = llm_prompt_service.create(data) + logger.info( + f"Successfully created llm prompt with ID: {new_prompt.id}" + ) + return new_prompt + except Exception as e: + logger.error(f"Error creating llm prompt: {str(e)}") + raise e + + +@router.put('/{prompt_id}') +async def update_prompt(prompt_id: int, file: LlmPromptSchema, + llm_prompt_service: LlmPromptService = Depends(DI.get_llm_prompt_service)): + logger.info("Handling PUT request to /llm_prompt/{prompt_id}") + try: + updated_prompt = llm_prompt_service.update(prompt_id, file) + logger.info( + f"Successfully updated llm prompt with ID: {prompt_id}" + ) + return updated_prompt + except Exception as e: + logger.error(f"Error updating llm prompt: {str(e)}") + raise e diff --git a/routes/log.py b/routes/log.py new file mode 100644 index 0000000000000000000000000000000000000000..16c3e9b1177cc72252e0be68e1bc40d1c62e82ad --- /dev/null +++ b/routes/log.py @@ -0,0 +1,116 @@ +import logging +from datetime import datetime +from typing import Annotated, Optional + +from fastapi import APIRouter, Depends, Query +from sqlalchemy.orm import aliased +from starlette import status + +from common.common import configure_logging +from common.exceptions import LogNotFoundException +from components.dbo.models.feedback import Feedback +from components.dbo.models.log import Log +from schemas.log import LogCreate +import common.dependencies as DI +from sqlalchemy.orm import sessionmaker + +router = APIRouter() + +logger = logging.getLogger(__name__) +configure_logging() + + +@router.get('/logs', status_code=status.HTTP_200_OK) +async def get_all_logs( + db: Annotated[sessionmaker, Depends(DI.get_db)], + date_start: Optional[datetime] = Query(None, alias="date_start"), + date_end: Optional[datetime] = Query(None, alias="date_end"), +): + logger.info(f'GET /logs: start') + logger.info(f'GET /logs: start_date={date_start}, end_date={date_end}') + feedback_alias = aliased(Feedback) + + query = db.query(Log) + + if date_start and date_end: + query = query.filter(Log.dateCreated.between(date_start, date_end)) + elif date_start: + query = query.filter(Log.dateCreated >= date_start) + elif date_end: + query = query.filter(Log.dateCreated <= date_end) + + query = query.outerjoin(feedback_alias, Log.id == feedback_alias.log_id) + + logs_with_feedback = query.all() + + combined_logs = [] + for log in logs_with_feedback: + if log.feedback: + for feedback in log.feedback: + combined_logs.append( + { + "log_id": log.id, + "llmPrompt": log.llmPrompt, + "llmResponse": log.llmResponse, + "llm_classifier": log.llm_classifier, + "dateCreated": log.dateCreated, + "userRequest": log.userRequest, + "userName": log.userName, + "query_type": log.query_type, + "feedback_id": feedback.feedback_id, + "userComment": feedback.userComment, + "userScore": feedback.userScore, + "manualEstimate": feedback.manualEstimate, + "llmEstimate": feedback.llmEstimate, + } + ) + else: + combined_logs.append( + { + "log_id": log.id, + "llmPrompt": log.llmPrompt, + "llmResponse": log.llmResponse, + "llm_classifier": log.llm_classifier, + "dateCreated": log.dateCreated, + "userRequest": log.userRequest, + "userName": log.userName, + "query_type": log.query_type, + "feedback_id": None, + "userComment": None, + "userScore": None, + "manualEstimate": None, + "llmEstimate": None, + } + ) + return combined_logs + + +@router.get('/log/{log_id}', status_code=status.HTTP_200_OK) +async def get_log(db: Annotated[sessionmaker, Depends(DI.get_db)], log_id): + log = db.query(Log).filter(Log.id == log_id).first() + if log is None: + raise LogNotFoundException(log_id) + return log + + +@router.post('/log', status_code=status.HTTP_201_CREATED) +async def create_log(log: LogCreate, db: Annotated[sessionmaker, Depends(DI.get_db)]): + logger.info("Handling POST request to /log") + try: + new_log = Log( + llmPrompt=log.llmPrompt, + llmResponse=log.llmResponse, + llm_classifier=log.llm_classifier, + userRequest=log.userRequest, + userName=log.userName, + ) + + db.add(new_log) + db.commit() + db.refresh(new_log) + + logger.info(f"Successfully created log with ID: {new_log.id}") + return new_log + except Exception as e: + logger.error(f"Error creating log: {str(e)}") + raise e diff --git a/schemas/__init__.py b/schemas/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/schemas/acronym.py b/schemas/acronym.py new file mode 100644 index 0000000000000000000000000000000000000000..55a841e0316c809a262359b1c02edce7e32b3e17 --- /dev/null +++ b/schemas/acronym.py @@ -0,0 +1,24 @@ +from datetime import datetime +from typing import Optional + +from pydantic import BaseModel + + +class AcronymCollectionResponse(BaseModel): + collection_id: int + collection_name: str + collection_filename: str + updated_at: Optional[datetime] + acronyms: dict[str, list[str]] + + class Config: + from_attributes = True + json_schema_extra = { + 'example': { + 'collection_id': 1, + 'collection_name': 'default', + 'collection_filename': '', + 'updated_at': '2024-01-01T00:00:00', + 'acronyms': {'AB': ['Alpha Beta', 'Alpha Beta Gamma']}, + } + } diff --git a/schemas/dataset.py b/schemas/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..4ad52ce6b18efc8cc468ea8bd940b7fd96090bc3 --- /dev/null +++ b/schemas/dataset.py @@ -0,0 +1,111 @@ +from datetime import datetime + +from pydantic import BaseModel + +from schemas.document import Document + + +class Dataset(BaseModel): + """ + Краткая информация о датасете. + """ + + id: int + dateCreated: datetime + name: str + isActive: bool + isDraft: bool + + class Config: + from_attributes = True + json_schema_extra = { + 'example': { + 'id': 1, + 'dateCreated': datetime.fromisoformat('2024-01-01T00:00:00'), + 'name': 'default', + 'isActive': True, + 'isDraft': False, + } + } + + +class DocumentsPage(BaseModel): + """ + Страница с документами. + """ + + page: list[Document] + total: int + pageNumber: int + pageSize: int + + class Config: + json_schema_extra = { + 'example': { + 'documents': [Document.Config.json_schema_extra['example']], + 'total': 500, + 'pageNumber': 1, + 'pageSize': 10, + } + } + + +class DatasetExpanded(Dataset): + """ + Расширенная информация о датасете, включающая страницу с документами. + """ + + data: DocumentsPage + + class Config: + json_schema_extra = { + 'example': { + 'id': 1, + 'dateCreated': datetime.fromisoformat('2024-01-01T00:00:00'), + 'name': 'default', + 'isActive': True, + 'isDraft': False, + 'data': DocumentsPage.Config.json_schema_extra['example'], + } + } + + +class DatasetProcessing(BaseModel): + """ + Информация о процессе обработки датасета. + """ + + status: str + total: int | None + current: int | None + datasetName: str | None + + class Config: + json_schema_extra = { + 'example': { + 'status': 'ready', + } + } + + +class SortQuery(BaseModel): + """ + Запрос на сортировку. + """ + + field: str + direction: str + + class Config: + json_schema_extra = {'example': {'field': 'name', 'direction': 'asc'}} + + +class SortQueryList(BaseModel): + """ + Список запросов на сортировку. + """ + + sorts: list[SortQuery] + + class Config: + json_schema_extra = {'example': [{'field': 'name', 'direction': 'asc'}]} diff --git a/schemas/document.py b/schemas/document.py new file mode 100644 index 0000000000000000000000000000000000000000..8b22f3b95070426af0462ef747a33bd22ece3fc3 --- /dev/null +++ b/schemas/document.py @@ -0,0 +1,33 @@ +import os + +from pydantic import BaseModel + + +class Document(BaseModel): + id: int + name: str | None + owner: str | None + status: str + + class Config: + json_schema_extra = { + 'example': { + 'id': 1, + 'name': 'Положение о положении', + 'owner': 'Отдел управления отделами', + 'status': 'Актуален', + } + } + + +class DocumentDownload(BaseModel): + filename: str + filepath: str | os.PathLike + + class Config: + json_schema_extra = { + 'example': { + 'filename': 'Положение о положении.xml', + 'filepath': '/path/to/file.xml', + } + } diff --git a/schemas/feedback.py b/schemas/feedback.py new file mode 100644 index 0000000000000000000000000000000000000000..3be260daa67f7a4e59007eddeae5591cb346a5a2 --- /dev/null +++ b/schemas/feedback.py @@ -0,0 +1,29 @@ +from pydantic import BaseModel, field_validator + +from exceptions import InvalidEstimateException, InvalidUserScoreException + + +class FeedbackCreate(BaseModel): + log_id: int + userComment: str + userScore: int + manualEstimate: int + llmEstimate: int + + @field_validator("userScore") + def check_user_score(cls, value): + if not (1 <= value <= 5): + raise InvalidUserScoreException(value) + return value + + @field_validator("manualEstimate") + def check_manual_estimate(cls, value): + if value < 1: + raise InvalidEstimateException(value) + return value + + @field_validator("llmEstimate") + def check_llm_estimate(cls, value): + if value < 1: + raise InvalidEstimateException(value) + return value diff --git a/schemas/llm_config.py b/schemas/llm_config.py new file mode 100644 index 0000000000000000000000000000000000000000..cde2e1952fcd33f1977ac1b32d33613c02bde7cd --- /dev/null +++ b/schemas/llm_config.py @@ -0,0 +1,57 @@ +from typing import Optional +from datetime import datetime +from pydantic import BaseModel + + +class LLMConfig(BaseModel): + is_default: bool + id: int + model: str + temperature: float + top_p: float + min_p: float + frequency_penalty: float + presence_penalty: float + n_predict: int + seed: int + date_created: datetime + + class Config: + json_schema_extra = { + 'example': { + 'is_default': True, + 'model': 'meta-llama/Llama-3.3-70B-Instruct', + 'temperature': 0.14, + 'top_p': 0.95, + 'min_p': 0.05, + 'frequency_penalty': -0.001, + 'presence_penalty': 1.3, + 'n_predict': 1000, + 'seed': 42 + } + } + +class LLMConfigCreateScheme(BaseModel): + is_default: bool + model: str + temperature: float + top_p: float + min_p: float + frequency_penalty: float + presence_penalty: float + n_predict: int + seed: int + class Config: + json_schema_extra = { + 'example': { + 'is_default': True, + 'model': 'meta-llama/Llama-3.3-70B-Instruct', + 'temperature': 0.14, + 'top_p': 0.95, + 'min_p': 0.05, + 'frequency_penalty': -0.001, + 'presence_penalty': 1.3, + 'n_predict': 1000, + 'seed': 42 + } + } diff --git a/schemas/llm_prompt.py b/schemas/llm_prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..89d174249f5ac4cb5cd1c12961a5bd0181dc56d2 --- /dev/null +++ b/schemas/llm_prompt.py @@ -0,0 +1,37 @@ +from typing import Optional +from datetime import datetime +from pydantic import BaseModel + + +class LlmPromptSchema(BaseModel): + is_default: bool + id: int + text: str + name: str + type: str + date_created: datetime + + class Config: + json_schema_extra = { + 'example': { + 'is_default': True, + 'text': 'Я большой и страшный промпт. Я в поросятах знаю толк.', + 'name': 'Большой и страшный промпт', + 'type': 'system' + } + } + +class LlmPromptCreateSchema(BaseModel): + is_default: bool + text: str + name: str + type: str + class Config: + json_schema_extra = { + 'example': { + 'is_default': True, + 'text': 'Я большой и страшный промпт. Я в поросятах знаю толк.', + 'name': 'Большой и страшный промпт', + 'type': 'system' + } + } diff --git a/schemas/log.py b/schemas/log.py new file mode 100644 index 0000000000000000000000000000000000000000..0935a28ac4bf3b6a195b3b8a2bcb2ea535b2a2a6 --- /dev/null +++ b/schemas/log.py @@ -0,0 +1,12 @@ +from typing import Optional + +from pydantic import BaseModel + + +class LogCreate(BaseModel): + llmPrompt: Optional[str] = None + llmResponse: Optional[str] = None + userRequest: str + llm_classifier: Optional[str] = None + query_type: Optional[str] = None + userName: str