update

business_transaction_map/components/faiss_vector_database.py

@@ -187,23 +187,46 @@ class FaissVectorDatabase:
                 'TypeDocs': К кому разделу относится карта проводок (1С или SAP)
             }
         """
-        if len(emb_query.shape) != 2:
-            assert print('Не правильный размер вектора!')
+         # Проверка размерности вектора
+        if len(emb_query.shape) != 2 or emb_query.shape[0] != 1:
+            raise ValueError("Вектор запроса должен иметь размерность [1, d].")
             
-        print("Index dimension:", self.index.d)  # Размерность индекса
-        print("Query dimension:", emb_query.shape[1])  # Размерность вектора запроса
-
-
-        distances, indexes = self.index.search(emb_query, k_neighbors)
+        # Приведение вектора запроса к нужной размерности, если требуется
+        query_dim = emb_query.shape[1]
+        index_dim = self.index.d
+        if query_dim != index_dim:
+            if query_dim < index_dim:
+                # Padding до нужной размерности
+                padded_query = torch.zeros((1, index_dim), dtype=emb_query.dtype)
+                padded_query[:, :query_dim] = emb_query
+                emb_query = padded_query
+            else:
+                raise ValueError(
+                    f"Размерность вектора запроса ({query_dim}) превышает размерность индекса ({index_dim})."
+                )
+            
+        emb_query_np = emb_query.numpy().astype("float32")
+
+        # Поиск в индексе
+        try:
+            distances, indexes = self.index.search(emb_query_np, k_neighbors)
+        except Exception as e:
+            raise RuntimeError(f"Ошибка при выполнении поиска: {e}")
+        
         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]
+            if ind < 0 or ind >= len(self.df):
+                answers[i] = {"error": "Индекс результата выходит за пределы таблицы"}
+                continue
+            
+            answers[i] = {
+                "distance": distances[0][i],
+                "index_answer": ind,
+                "doc_name": self.df.iloc[ind][COLUMN_DOC_NAME],
+                "text_answer": self.df.iloc[ind][COLUMN_TABLE_NAME],
+                COLUMN_LABELS_STR: self.df.iloc[ind][COLUMN_LABELS_STR],
+                COLUMN_NAMES: self.df.iloc[ind][COLUMN_NAMES],
+                COLUMN_TYPE_DOC_MAP: self.df.iloc[ind][COLUMN_TYPE_DOC_MAP],
+            }
 
         return answers