Update: new semantic search logic was provided, requirements.txt file was fixed

app.py

@@ -5,7 +5,6 @@ from semantic_search import SemanticSearch
 from datetime import datetime
 
 search = SemanticSearch()
-search.test_search()
 
 app = Flask(__name__)
 app.config['JSON_AS_ASCII'] = False
@@ -20,6 +19,7 @@ if not os.path.exists(LOGS_BASE_PATH):
 # Check if logs are enabled
 ENABLE_LOGS = os.getenv("ENABLE_LOGS", "0") == "1"
 
+
 def log_query_result(query, top, request_id, result):
     if not ENABLE_LOGS:
         return
@@ -38,10 +38,12 @@ def log_query_result(query, top, request_id, result):
     with open(log_file_path, 'w') as log_file:
         json.dump(log_data, log_file, indent=2)
 
+
 @app.route('/health', methods=['GET'])
 def health():
     return jsonify({"status": "ok"})
 
+
 @app.route('/search', methods=['POST'])
 def search_route():
     data = request.get_json()
@@ -56,6 +58,7 @@ def search_route():
 
     return jsonify(result)
 
+
 @app.route('/read_logs', methods=['GET'])
 def read_logs():
     logs = []
@@ -66,6 +69,7 @@ def read_logs():
                 logs.append(log_data)
     return jsonify(logs)
 
+
 @app.route('/analyze_logs', methods=['GET'])
 def analyze_logs():
     logs_by_query_top = {}
@@ -91,5 +95,6 @@ def analyze_logs():
 
     return jsonify(invalid_logs)
 
+
 if __name__ == '__main__':
     app.run(debug=False, host='0.0.0.0')

requirements.txt

@@ -8,4 +8,5 @@ transformers==4.29.2
 # sentencepiece==0.1.99
 # six==1.16.0
 # tokenizers==0.13.3
-flask==3.0.0
+flask==3.0.0
+datasets

semantic_search.py

@@ -1,54 +1,58 @@
 import os
-import json
+# import json
 import torch
-import pickle
+# import pickle
 import numpy as np
 import faiss
-
+from datasets import Dataset as dataset
 from transformers import AutoTokenizer, AutoModel
-from legal_info_search_utils.utils import get_subsets_for_db, get_subsets_for_qa
-from legal_info_search_utils.utils import filter_db_data_types, filter_qa_data_types
-from legal_info_search_utils.utils import db_tokenization, qa_tokenization
-from legal_info_search_utils.utils import extract_text_embeddings, filter_ref_parts
-from legal_info_search_utils.utils import print_metrics, get_final_metrics
-from legal_info_search_utils.utils import court_text_splitter
+# from legal_info_search_utils.utils import get_subsets_for_db, get_subsets_for_qa
+# from legal_info_search_utils.utils import filter_db_data_types, filter_qa_data_types
+# from legal_info_search_utils.utils import db_tokenization, qa_tokenization
+# from legal_info_search_utils.utils import extract_text_embeddings, filter_ref_parts
+# from legal_info_search_utils.utils import print_metrics, get_final_metrics
+# from legal_info_search_utils.utils import court_text_splitter
 from legal_info_search_utils.utils import query_tokenization, query_embed_extraction
-from legal_info_search_utils.metrics import calculate_metrics_at_k
+
+# from legal_info_search_utils.metrics import calculate_metrics_at_k
 
 
 global_data_path = os.environ.get("GLOBAL_DATA_PATH", "legal_info_search_data/")
-global_model_path = os.environ.get("GLOBAL_MODEL_PATH", 
+global_model_path = os.environ.get("GLOBAL_MODEL_PATH",
                                    "legal_info_search_model/20240202_204910_ep8/")
 
 # размеченные консультации
-data_path_consult = os.environ.get("DATA_PATH_CONSULT", 
-                                   global_data_path + "data_jsons_20240202.pkl")
+# data_path_consult = os.environ.get("DATA_PATH_CONSULT",
+#                                    global_data_path + "data_jsons_20240202.pkl")
+
+data_path_consult = os.environ.get("DATA_PATH_CONSULT",
+                                   global_data_path + "court_dataset_chunk_200_correct_tokenizer_for_develop")
 
 # id консультаций, побитые на train / valid / test
-data_path_consult_ids = os.environ.get("DATA_PATH_CONSULT_IDS",
-                                       global_data_path + "data_ids.json")
+# data_path_consult_ids = os.environ.get("DATA_PATH_CONSULT_IDS",
+#                                        global_data_path + "data_ids.json")
 
 # предобработанные внутренние документы
-data_path_internal_docs = os.environ.get("DATA_PATH_INTERNAL_DOCS",
-                                       global_data_path + "internal_docs.json")
+# data_path_internal_docs = os.environ.get("DATA_PATH_INTERNAL_DOCS",
+#                                        global_data_path + "internal_docs.json")
 
 # состав БД
 # $ export DB_SUBSETS='["train", "valid", "test"]'
-db_subsets = os.environ.get("DB_SUBSETS", ["train", "valid", "test"])
+# db_subsets = os.environ.get("DB_SUBSETS", ["train", "valid", "test"])
 
 # Отбор типов документов. В списке указать те, которые нужно оставить в БД.
 # $ export DB_DATA_TYPES='["НКРФ", "ГКРФ", "ТКРФ"]'
 db_data_types = os.environ.get("DB_DATA_TYPES", [
-    'НКРФ',
-    'ГКРФ',
-    'ТКРФ',
-    'Федеральный закон',
+    # 'НКРФ',
+    # 'ГКРФ',
+    # 'ТКРФ',
+    # 'Федеральный закон',
     'Письмо Минфина',
     'Письмо ФНС',
-    'Приказ ФНС',
-    'Постановление Правительства',
+    # 'Приказ ФНС',
+    # 'Постановление Правительства',
     'Судебный документ',
-    'Внутренний документ'
+    # 'Внутренний документ'
 ])
 
 device = os.environ.get("MODEL_DEVICE", 'cuda' if torch.cuda.is_available() else 'cpu')
@@ -57,67 +61,22 @@ device = os.environ.get("MODEL_DEVICE", 'cuda' if torch.cuda.is_available() else
 hf_token = os.environ.get("HF_TOKEN", "")
 hf_model_name = os.environ.get("HF_MODEL_NAME", "")
 
+
 class SemanticSearch:
-    def __init__(self, index_type="IndexFlatIP", do_embedding_norm=True, 
-                 faiss_batch_size=8, do_normalization=True):
+    def __init__(self,
+                 index_type="IndexFlatIP",
+                 do_embedding_norm=True,
+                 do_normalization=True):
         self.device = device
         self.do_embedding_norm = do_embedding_norm
-        self.faiss_batch_size = faiss_batch_size
         self.do_normalization = do_normalization
         self.load_model()
-        
+
         indexes = {
             "IndexFlatL2": faiss.IndexFlatL2(self.embedding_dim),
             "IndexFlatIP": faiss.IndexFlatIP(self.embedding_dim)
         }
-        self.index = indexes[index_type]
-        self.load_data()
-        self.preproces_data()
-        self.test_search()
 
-    def load_data(self):
-        with open(data_path_consult, "rb") as f:
-            all_docs = pickle.load(f)
-
-        with open(data_path_consult_ids, "r", encoding="utf-8") as f:
-            data_ids = json.load(f)
-
-        with open(data_path_internal_docs, "r", encoding="utf-8") as f:
-            internal_docs = json.load(f)
-
-        db_data = get_subsets_for_db(db_subsets, data_ids, all_docs)
-        filtered_all_docs = filter_qa_data_types(db_data_types, all_docs)
-        
-        self.mean_refs_count = self.get_mean_refs_counts(db_data_types, filtered_all_docs)
-        self.mean_refs_count['Внутренний документ'] = 3
-        self.filtered_db_data = filter_db_data_types(db_data_types, db_data)
-        self.filtered_db_data.update(internal_docs)
-        self.all_docs_qa = get_subsets_for_qa(["valid"], data_ids, filtered_all_docs)
-
-    def load_model(self):
-        if hf_token and hf_model_name:
-            self.tokenizer = AutoTokenizer.from_pretrained(hf_model_name, use_auth_token=True)
-            self.model = AutoModel.from_pretrained(hf_model_name, use_auth_token=True).to(self.device)
-        else:
-            self.tokenizer = AutoTokenizer.from_pretrained(global_model_path)
-            self.model = AutoModel.from_pretrained(global_model_path).to(self.device)
-            
-        self.max_len = self.tokenizer.max_len_single_sentence
-        self.embedding_dim = self.model.config.hidden_size
-
-    def preproces_data(self):
-        index_keys, index_toks = db_tokenization(self.filtered_db_data, self.tokenizer)
-        val_questions, val_refs = qa_tokenization(self.all_docs_qa, self.tokenizer)
-        docs_embeds_faiss, questions_embeds_faiss = extract_text_embeddings(index_toks, 
-            val_questions, self.model, self.do_normalization, self.faiss_batch_size)
-
-        self.index.add(docs_embeds_faiss)
-        self.index_keys = index_keys
-        self.index_toks = index_toks
-        self.val_questions = val_questions
-        self.val_refs = val_refs
-        self.docs_embeds_faiss = docs_embeds_faiss
-        self.questions_embeds_faiss = questions_embeds_faiss
         self.optimal_params = {
             'НКРФ': {'thresh': 0.58421, 'sim_factor': 0.89474},
             'ГКРФ': {'thresh': 0.64737, 'sim_factor': 0.89474},
@@ -130,61 +89,33 @@ class SemanticSearch:
             'Судебный документ': {'thresh': 0.67895, 'sim_factor': 0.89474},
             'Внутренний документ': {'thresh': 0.55263, 'sim_factor': 0.84211}
         }
-        self.ref_categories = {
-            'all': 'all',
-            'НКРФ': 'НКРФ',
-            'ГКРФ': 'ГКРФ',
-            'ТКРФ': 'ТКРФ',
-            'Федеральный закон': 'ФЗ',
-            'Судебный документ': 'Суды',
-            'Письмо Минфина': 'Письмо МФ',
-            'Письмо ФНС': 'Письмо ФНС',
-            'Приказ ФНС': 'Приказ ФНС',
-            'Постановление Правительства': 'Пост. Прав.',
-            'Внутренний документ': 'Внутр. док.'
-        }
-
-    def test_search(self):
-        topk = len(self.filtered_db_data)
-        pred_raw = {}
-        true = {}
-        all_distances = []
-        for idx, (q_embed, refs) in enumerate(zip(self.questions_embeds_faiss, 
-                                                  self.val_refs.values())):
-            distances, indices = self.index.search(np.expand_dims(q_embed, 0), topk)
-            pred_raw[idx] = [self.index_keys[x] for x in indices[0]]
-            true[idx] = list(refs)
-            all_distances.append(distances)
+        self.index_type = index_type
+        self.index_docs = indexes[self.index_type]
+        self.load_data()
+        self.docs_embeddings = [torch.unsqueeze(torch.Tensor(x['doc_embedding']), 0) for x in self.all_docs_info]
+        self.docs_embeddings = torch.cat(self.docs_embeddings, dim=0)
+        self.index_docs.add(self.docs_embeddings)
+        # self.preproces_data()
+        # self.test_search()
 
-        pred = {}
-        for idx, p, d in zip(true.keys(), pred_raw.values(), all_distances):
-            fp, fs = self.search_results_filtering(p, d[0])
-            pred[idx] = fp
+    def load_data(self):
+        self.all_docs_info = dataset.load_from_disk(data_path_consult).to_list()
+        self.docs_names = [doc['doc_name'] for doc in self.all_docs_info]
+        self.mean_refs_count = {'Письмо Минфина': 3,
+                                'Письмо ФНС': 2,
+                                'Судебный документ': 3}
 
-        # раскомментировать нужное. Если всё закомментировано - метрики
-        # посчтаются "как есть", с учетом полной иерархии
-        filter_parts = [
-        #     "абз.", 
-        #     "пп.", 
-        #     "п."
-        ]
-        filtered_pred = filter_ref_parts(pred, filter_parts)
-        filtered_true = filter_ref_parts(true, filter_parts)
+    def load_model(self):
+        if hf_token and hf_model_name:
+            self.tokenizer = AutoTokenizer.from_pretrained(hf_model_name, use_auth_token=True)
+            self.model = AutoModel.from_pretrained(hf_model_name, use_auth_token=True).to(self.device)
+        else:
+            self.tokenizer = AutoTokenizer.from_pretrained(global_model_path)
+            self.model = AutoModel.from_pretrained(global_model_path).to(self.device)
 
-        metrics_func_params = {
-            'compensate_div_0': True, 
-            'dynamic_topk': True,
-            'skip_empty_trues': False,
-            'skip_empty_preds': False
-        }
-        metrics = get_final_metrics(filtered_pred, filtered_true, 
-            self.ref_categories.keys(), [0], 
-            metrics_func=calculate_metrics_at_k, 
-            metrics_func_params=metrics_func_params)
-        
-        print_metrics(metrics, self.ref_categories)
+        self.max_len = self.tokenizer.max_len_single_sentence
+        self.embedding_dim = self.model.config.hidden_size
 
-        
     def search_results_filtering(self, pred, dists):
         all_ctg_preds = []
         all_scores = []
@@ -192,7 +123,7 @@ class SemanticSearch:
             ctg_thresh = self.optimal_params[ctg]['thresh']
             ctg_sim_factor = self.optimal_params[ctg]['sim_factor']
 
-            ctg_preds = [(ref, dist) for ref, dist in zip(pred, dists) 
+            ctg_preds = [(ref, dist) for ref, dist in zip(pred, dists)
                          if ctg in ref and dist > ctg_thresh]
 
             sorted_pd = sorted(ctg_preds, key=lambda x: x[1], reverse=True)
@@ -222,48 +153,39 @@ class SemanticSearch:
 
         return sorted_preds, sorted_scores
 
-    def court_docs_shrinking(self, preds, docs):
-        new_preds = []
-        new_docs = []
-
-        for ref_name, ref_text in zip(preds, docs):
-            is_court = 'Судебный документ' in ref_name
-            has_splitter = court_text_splitter in ref_text
-
-            if is_court and has_splitter:
-                new_ref_text = ref_text.split(court_text_splitter)[0].strip()
-                new_preds.append(ref_name)
-                new_docs.append(new_ref_text)
-            else:
-                new_preds.append(ref_name)
-                new_docs.append(ref_text)
-        return new_preds, new_docs
-
-    @staticmethod
-    def get_mean_refs_counts(db_data_types, data):
-        mean_refs_count = {}
-        for tp in db_data_types:
-            all_tp_refs = []
-            for doc in data.values():
-                tp_refs_len = len([ref for ref in doc['added_refs'] if tp in ref])
-                if tp_refs_len:
-                    all_tp_refs.append(tp_refs_len)
-            
-            mean_refs_count[tp] = np.mean(all_tp_refs)
-
-        for k, v in mean_refs_count.items():
-            mean_refs_count[k] = int(v + 1)
-
-        return mean_refs_count
+    def get_most_relevant_teaser(self,
+                                 question: str = None,
+                                 doc_index: int = None):
+        teaser_indexes = {
+            "IndexFlatL2": faiss.IndexFlatL2(self.embedding_dim),
+            "IndexFlatIP": faiss.IndexFlatIP(self.embedding_dim)
+        }
+        teasers_index = teaser_indexes[self.index_type]
+        question_tokens = query_tokenization(question, self.tokenizer)
+        question_embedding = query_embed_extraction(question_tokens, self.model,
+                                                    self.do_normalization)
+        # question_embedding = self.custom_embedder.create_question_embedding(question)
+        teasers_texts = [teaser['summary_text'] for teaser in self.all_docs_info[doc_index]['chunks_embeddings']]
+        teasers_embeddings = [torch.unsqueeze(torch.Tensor(teaser['embedding']), 0) for teaser in
+                              self.all_docs_info[doc_index]['chunks_embeddings']]
+        teasers_embeddings = torch.cat(teasers_embeddings, 0)
+        teasers_index.add(teasers_embeddings)
+        distances, indices = teasers_index.search(question_embedding, 10)
+        most_relevant_teaser = teasers_texts[indices[0][0]]
+        return most_relevant_teaser
 
     def search(self, query, top=15):
         query_tokens = query_tokenization(query, self.tokenizer)
-        query_embeds = query_embed_extraction(query_tokens, self.model, 
+        query_embeds = query_embed_extraction(query_tokens, self.model,
                                               self.do_normalization)
-        distances, indices = self.index.search(query_embeds, len(self.filtered_db_data))
-        pred = [self.index_keys[x] for x in indices[0]]
+        distances, indices = self.index_docs.search(query_embeds, len(self.all_docs_info))
+        pred = [self.all_docs_info[x]['doc_name'] for x in indices[0]]
         preds, scores = self.search_results_filtering(pred, distances[0])
-        docs = [self.filtered_db_data[ref] for ref in preds]
-        preds, docs = self.court_docs_shrinking(preds, docs)
-
+        # docs = [self.all_docs_info[x][ref] for ref in preds]
+        docs = []
+        for ref in preds:
+            doc_index = self.docs_names.index(ref)
+            most_relevant_teaser = self.get_most_relevant_teaser(question=query,
+                                                                 doc_index=doc_index)
+            docs.append(most_relevant_teaser)
         return preds[:top], docs[:top], scores[:top]