Search updated

app.py

@@ -1,5 +1,4 @@
 import json
-import os
 from flask import Flask, jsonify, request
 from semantic_search import SemanticSearch
 
@@ -24,4 +23,4 @@ def search_route():
 
 if __name__ == '__main__':
     
-    app.run(debug=False, host='0.0.0.0', port=7868)
+    app.run(debug=False, host='0.0.0.0')

legal_info_search_data/data_jsons_20240202.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:78cab704acf861b87eec01ba4d575e2e0110ed57ac64f814c04c3de02ef2db88
+size 22359347

legal_info_search_utils/metrics.py

@@ -1,7 +1,8 @@
 import numpy as np
 
 
-def calculate_metrics_at_k(pred, true, k, dynamic_topk=False):
+def calculate_metrics_at_k(pred, true, k, compensate_div_0=False, dynamic_topk=True, 
+                           skip_empty_trues=False, skip_empty_preds=False):
     precisions_at_k = []
     recalls_at_k = []
     f1_scores_at_k = []
@@ -14,11 +15,17 @@ def calculate_metrics_at_k(pred, true, k, dynamic_topk=False):
         relevant_documents = set(true[query_id])
         true_positives = len(retrieved_documents.intersection(relevant_documents))
 
-        if not len(retrieved_documents) and not len(relevant_documents):
+        if compensate_div_0 and not len(retrieved_documents) and not len(relevant_documents):
             precisions_at_k.append(1)
             recalls_at_k.append(1)
             f1_scores_at_k.append(1)
             continue
+
+        if skip_empty_trues and not len(relevant_documents):
+            continue
+
+        if skip_empty_preds and not len(retrieved_documents):
+            continue
         
         # precision
         precision_at_k = true_positives / k if k else 0

legal_info_search_utils/utils.py

@@ -10,8 +10,6 @@ from torch.utils.data import Dataset, DataLoader
 from torch.cuda.amp import autocast
 
 
-all_types_but_courts = ['НКРФ', 'ГКРФ', 'ТКРФ', 'Федеральный закон', 'Письмо Минфина', 
-                        'Письмо ФНС', 'Приказ ФНС', 'Постановление Правительства']
 court_text_splitter = "Весь текст судебного документа: "
 
 
@@ -54,12 +52,8 @@ def get_subsets_for_qa(subsets, data_ids, all_docs):
 def filter_db_data_types(text_parts, db_data_in):
     filtered_db_data = {}
     db_data = copy.deepcopy(db_data_in)
-    check_if_courts = 'Суды' in text_parts
-    for ref, text in db_data.items():
-        check_not_other = not any([True for x in all_types_but_courts if x in ref])
-        court_condition = check_if_courts and check_not_other
-        
-        if court_condition or any([True for x in text_parts if x in ref]):
+    for ref, text in db_data.items():        
+        if any([True for x in text_parts if x in ref]):
             filtered_db_data[ref] = text
     return filtered_db_data
 
@@ -73,12 +67,8 @@ def filter_qa_data_types(text_parts, all_docs_in):
             continue
 
         filtered_refs = {}
-        check_if_courts = 'Суды' in text_parts
         for ref, text in doc['added_refs'].items():
-            check_not_other = not any([True for x in all_types_but_courts if x in ref])
-            court_condition = check_if_courts and check_not_other
-            
-            if court_condition or any([True for x in text_parts if x in ref]):
+            if any([True for x in text_parts if x in ref]):
                 filtered_refs[ref] = text
         
         filtered_all_docs[doc_key] = doc
@@ -205,16 +195,9 @@ def get_exact_ctg_data(pred_in, true_in, ctg):
     
     out_pred = {}
     out_true = {}
-    check_if_courts = ctg == "Суды"
     for idx, (pred, true) in zip(true_in.keys(), zip(pred_in.values(), true_in.values())):
-        if check_if_courts:
-            ctg_refs_true = [ref for ref in true 
-                             if not any([True for x in all_types_but_courts if x in ref])]
-            ctg_refs_pred = [ref for ref in pred 
-                             if not any([True for x in all_types_but_courts if x in ref])]
-        else:
-            ctg_refs_true = [ref for ref in true if ctg in ref]
-            ctg_refs_pred = [ref for ref in pred if ctg in ref]
+        ctg_refs_true = [ref for ref in true if ctg in ref]
+        ctg_refs_pred = [ref for ref in pred if ctg in ref]
         
         out_true[idx] = ctg_refs_true
         out_pred[idx] = ctg_refs_pred

semantic_search.py

@@ -11,7 +11,7 @@ from legal_info_search_utils.utils import filter_db_data_types, filter_qa_data_t
 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 all_types_but_courts, court_text_splitter
+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
 
@@ -22,7 +22,7 @@ global_model_path = os.environ.get("GLOBAL_MODEL_PATH",
 
 # размеченные консультации
 data_path_consult = os.environ.get("DATA_PATH_CONSULT", 
-                                   global_data_path + "data_jsons_20240131.pkl")
+                                   global_data_path + "data_jsons_20240202.pkl")
 
 # id консультаций, побитые на train / valid / test
 data_path_consult_ids = os.environ.get("DATA_PATH_CONSULT_IDS",
@@ -43,7 +43,8 @@ db_data_types = os.environ.get("DB_DATA_TYPES", [
     'Письмо ФНС',
     'Приказ ФНС',
     'Постановление Правительства',
-    'Суды'
+    'Судебный документ',
+    'Внутренний документ'
 ])
 
 device = os.environ.get("MODEL_DEVICE", 'cuda' if torch.cuda.is_available() else 'cpu')
@@ -80,6 +81,7 @@ class SemanticSearch:
         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.filtered_db_data = filter_db_data_types(db_data_types, db_data)
         self.all_docs_qa = get_subsets_for_qa(["valid"], data_ids, filtered_all_docs)
 
@@ -108,24 +110,16 @@ class SemanticSearch:
         self.docs_embeds_faiss = docs_embeds_faiss
         self.questions_embeds_faiss = questions_embeds_faiss
         self.optimal_params = {
-            'НКРФ': {
-            'thresh': 0.613793, 'sim_factor': 0.878947, 'diff_n': 0},
-            'ГКРФ': {
-            'thresh': 0.758620, 'sim_factor': 0.878947, 'diff_n': 0},
-            'ТКРФ': {
-            'thresh': 0.734482, 'sim_factor': 0.9, 'diff_n': 0},
-            'Федеральный закон': {
-            'thresh': 0.734482, 'sim_factor': 0.5, 'diff_n': 0},
-            'Письмо Минфина': {
-            'thresh': 0.782758, 'sim_factor': 0.5, 'diff_n': 0},
-            'Письмо ФНС': {
-            'thresh': 0.879310, 'sim_factor': 0.5, 'diff_n': 0},
-            'Приказ ФНС': {
-            'thresh': 0.806896, 'sim_factor': 0.5, 'diff_n': 0},
-            'Постановление Правительства': {
-            'thresh': 0.782758, 'sim_factor': 0.5, 'diff_n': 0},
-            'Суды': {
-            'thresh': 0.846153, 'sim_factor': 0.939230,'diff_n': 0}
+            'НКРФ': {'thresh': 0.61579, 'sim_factor': 0.84211},
+            'ГКРФ': {'thresh': 0.55263, 'sim_factor': 0.0},
+            'ТКРФ': {'thresh': 0.48947, 'sim_factor': 1.0},
+            'Федеральный закон': {'thresh': 0.52105, 'sim_factor': 0.94737},
+            'Письмо Минфина': {'thresh': 0.71053, 'sim_factor': 0.0},
+            'Письмо ФНС': {'thresh': 0.61579, 'sim_factor': 0.84211},
+            'Приказ ФНС': {'thresh': 0.52105, 'sim_factor': 0.94737},
+            'Постановление Правительства': {'thresh': 0.45789, 'sim_factor': 0.89474},
+            'Судебный документ': {'thresh': 0.80526, 'sim_factor': 0.89474},
+            'Внутренний документ': {'thresh': 0.71053, 'sim_factor': 0.0}
         }
         self.ref_categories = {
             'all': 'all',
@@ -133,11 +127,12 @@ class SemanticSearch:
             'ГКРФ': 'ГКРФ',
             'ТКРФ': 'ТКРФ',
             'Федеральный закон': 'ФЗ',
-            'Суды': 'Суды',
+            'Судебный документ': 'Суды',
             'Письмо Минфина': 'Письмо МФ',
             'Письмо ФНС': 'Письмо ФНС',
             'Приказ ФНС': 'Приказ ФНС',
-            'Постановление Правительства': 'Пост. Прав.'
+            'Постановление Правительства': 'Пост. Прав.',
+            'Внутренний документ': 'Внутр. док.'
         }
 
     def test_search(self):
@@ -157,7 +152,7 @@ class SemanticSearch:
             fp, fs = self.search_results_filtering(p, d[0])
             pred[idx] = fp
 
-        # раскомментировать нужное. Если всё закомментировано - метрики 
+        # раскомментировать нужное. Если в��ё закомментировано - метрики
         # посчтаются "как есть", с учетом полной иерархии
         filter_parts = [
         #     "абз.", 
@@ -168,8 +163,10 @@ class SemanticSearch:
         filtered_true = filter_ref_parts(true, filter_parts)
 
         metrics_func_params = {
-            # 'compensate_div_0': True, 
-            'dynamic_topk': True
+            '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], 
@@ -185,26 +182,15 @@ class SemanticSearch:
         for ctg in db_data_types:
             ctg_thresh = self.optimal_params[ctg]['thresh']
             ctg_sim_factor = self.optimal_params[ctg]['sim_factor']
-            ctg_diff_n = self.optimal_params[ctg]['diff_n']
 
-            if ctg == 'Суды':
-                ctg_preds = [(ref, dist) for ref, dist in zip(pred, dists) if dist > ctg_thresh 
-                             and not any([True for type_ in all_types_but_courts if type_ in ref])]
-            else:
-                ctg_preds = [(ref, dist) for ref, dist in zip(pred, dists) 
-                             if ctg in ref and dist > ctg_thresh]
+            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)
             sorted_preds = [x[0] for x in sorted_pd]
             sorted_dists = [x[1] for x in sorted_pd]
 
             if len(sorted_dists):
-                diffs = np.diff(sorted_dists, ctg_diff_n)
-                if len(diffs):
-                    n_preds = np.argmax(diffs) + ctg_diff_n + 1
-                else:
-                    n_preds = 0
-
                 if len(sorted_dists) > 1:
                     ratios = (sorted_dists[1:] / sorted_dists[0]) >= ctg_sim_factor
                     ratios = np.array([True, *ratios])
@@ -213,15 +199,12 @@ class SemanticSearch:
 
                 main_preds = np.array(sorted_preds)[np.where(ratios)].tolist()
                 scores = np.array(sorted_dists)[np.where(ratios)].tolist()
-                if ctg_diff_n > 0 and n_preds > 0:
-                    main_preds = main_preds[:n_preds]
-                    scores = scores[:n_preds]
             else:
                 main_preds = []
                 scores = []
 
-            all_ctg_preds.extend(main_preds)
-            all_scores.extend(scores)
+            all_ctg_preds.extend(main_preds[:self.mean_refs_count[ctg]])
+            all_scores.extend(scores[:self.mean_refs_count[ctg]])
 
         sorted_values = [(ref, score) for ref, score in zip(all_ctg_preds, all_scores)]
         sorted_values = sorted(sorted_values, key=lambda x: x[1], reverse=True)
@@ -235,7 +218,7 @@ class SemanticSearch:
         new_docs = []
 
         for ref_name, ref_text in zip(preds, docs):
-            is_court = not any([True for type_ in all_types_but_courts if type_ in ref_name])
+            is_court = 'Судебный документ' in ref_name
             has_splitter = court_text_splitter in ref_text
 
             if is_court and has_splitter:
@@ -247,7 +230,24 @@ class SemanticSearch:
                 new_docs.append(ref_text)
         return new_preds, new_docs
 
-    def search(self, query, top=10):
+    @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 search(self, query, top=15):
         query_tokens = query_tokenization(query, self.tokenizer)
         query_embeds = query_embed_extraction(query_tokens, self.model, 
                                               self.do_normalization)