Spaces:
Sleeping
Sleeping
File size: 9,621 Bytes
b24d496 bfc27c5 b24d496 bfc27c5 b24d496 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 |
from typing import Tuple, List, Dict, Union
import faiss
import pandas as pd
import numpy as np
import torch
from business_transaction_map.common.constants import COLUMN_DOC_NAME
from business_transaction_map.common.constants import COLUMN_EMBEDDING
from business_transaction_map.common.constants import COLUMN_LABELS_STR
from business_transaction_map.common.constants import COLUMN_NAMES
from business_transaction_map.common.constants import COLUMN_TABLE_NAME
from business_transaction_map.common.constants import COLUMN_TYPE_DOC_MAP
class FaissVectorDatabase:
"""Класс для взаимодействия между векторами и информацией о них"""
def __init__(self, path_to_metadata: str):
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;
"""
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']]
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.iloc[0]["Text"]}')
return other_info, ind
else:
if self.df.iloc[ind]['Pargaraph'] is None:
pass
else:
pattern = self.df.iloc[ind]["Pargaraph"].replace(".", r"\.")
paragraph, start_index_paragraph = self._paragraph_content(fr'^{pattern}?$', doc_number, ind, 1)
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)
else:
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'] = distances[0][i]
answers[i][f'index_answer'] = ind
answers[i][f'doc_name'] = self.df.iloc[ind]['DocName']
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
|