updated search

legal_info_search_utils/utils.py

@@ -12,6 +12,7 @@ from torch.cuda.amp import autocast
 
 all_types_but_courts = ['НКРФ', 'ГКРФ', 'ТКРФ', 'Федеральный закон', 'Письмо Минфина', 
                         'Письмо ФНС', 'Приказ ФНС', 'Постановление Правительства']
+court_text_splitter = "Весь текст судебного документа: "
 
 
 class FaissDocsDataset(Dataset):
@@ -172,26 +173,6 @@ def extract_text_embeddings(index_toks, val_questions, model,
     return docs_embeds_faiss.numpy(), questions_embeds_faiss.numpy()
 
 
-def run_semantic_search(index, model, tokenizer, filtered_db_data, all_docs_qa, 
-                        do_normalization=True, faiss_batch_size=16, topk=100):
-    index_keys, index_toks = db_tokenization(filtered_db_data, tokenizer)
-    val_questions, val_refs = qa_tokenization(all_docs_qa, tokenizer)
-    docs_embeds_faiss, questions_embeds_faiss = extract_text_embeddings(index_toks, 
-        val_questions, model, do_normalization, faiss_batch_size)
-    index.add(docs_embeds_faiss)
-
-    pred = {}
-    true = {}
-    all_distances = []
-    for idx, (q_embed, refs) in enumerate(zip(questions_embeds_faiss, val_refs.values())):
-        distances, indices = index.search(np.expand_dims(q_embed, 0), topk)
-        pred[idx] = [index_keys[x] for x in indices[0]]
-        true[idx] = list(refs)
-        all_distances.append(distances)
-
-    return pred, true, all_distances
-
-
 def filter_ref_parts(ref_dict, filter_parts):
     filtered_dict = {}
     for k, refs in ref_dict.items():
@@ -203,13 +184,13 @@ def filter_ref_parts(ref_dict, filter_parts):
 
 
 def get_final_metrics(pred, true, categories, top_k_values, 
-                      metrics_func, dynamic_topk=False):
+                      metrics_func, metrics_func_params):
     metrics = {}
     for top_k in top_k_values:
         ctg_metrics = {}
         for ctg in categories:
             ctg_pred, ctg_true = get_exact_ctg_data(pred, true, ctg)
-            metrics_at_k = metrics_func(ctg_pred, ctg_true, top_k, dynamic_topk)
+            metrics_at_k = metrics_func(ctg_pred, ctg_true, top_k, **metrics_func_params)
             for mk in metrics_at_k.keys():
                 metrics_at_k[mk] = round(metrics_at_k[mk] * 100, 6)
             ctg_metrics[ctg] = metrics_at_k

semantic_search.py

@@ -11,16 +11,18 @@ 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 query_tokenization, query_embed_extraction
 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", "e5_large_rus_finetuned_20240120_122822_ep6")
+global_model_path = os.environ.get("GLOBAL_MODEL_PATH", 
+                                   "legal_info_search_model/20240120_122822_ep6/")
 
 # размеченные консультации
 data_path_consult = os.environ.get("DATA_PATH_CONSULT", 
-                                   global_data_path + "data_jsons_20240119.pkl")
+                                   global_data_path + "data_jsons_20240131.pkl")
 
 # id консультаций, побитые на train / valid / test
 data_path_consult_ids = os.environ.get("DATA_PATH_CONSULT_IDS",
@@ -40,7 +42,8 @@ db_data_types = os.environ.get("DB_DATA_TYPES", [
     'Письмо Минфина',
     'Письмо ФНС',
     'Приказ ФНС',
-    'Постановление Правительства'
+    'Постановление Правительства',
+    'Суды'
 ])
 
 device = os.environ.get("MODEL_DEVICE", 'cuda' if torch.cuda.is_available() else 'cpu')
@@ -120,7 +123,9 @@ class SemanticSearch:
             'Приказ ФНС': {
             'thresh': 0.806896, 'sim_factor': 0.5, 'diff_n': 0},
             'Постановление Правительства': {
-            'thresh': 0.782758, '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}
         }
         self.ref_categories = {
             'all': 'all',
@@ -162,9 +167,14 @@ class SemanticSearch:
         filtered_pred = filter_ref_parts(pred, filter_parts)
         filtered_true = filter_ref_parts(true, filter_parts)
 
+        metrics_func_params = {
+            # 'compensate_div_0': True, 
+            'dynamic_topk': True
+        }
         metrics = get_final_metrics(filtered_pred, filtered_true, 
             self.ref_categories.keys(), [0], 
-            metrics_func=calculate_metrics_at_k, dynamic_topk=True)
+            metrics_func=calculate_metrics_at_k, 
+            metrics_func_params=metrics_func_params)
         
         print_metrics(metrics, self.ref_categories)
 
@@ -176,9 +186,13 @@ class SemanticSearch:
             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']
-            
-            ctg_preds = [(ref, dist) for ref, dist in zip(pred, dists) 
-                         if ctg in ref and dist > ctg_thresh]
+
+            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]
 
             sorted_pd = sorted(ctg_preds, key=lambda x: x[1], reverse=True)
             sorted_preds = [x[0] for x in sorted_pd]
@@ -216,6 +230,23 @@ 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 = not any([True for type_ in all_types_but_courts if type_ 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
+
     def search(self, query, top=10):
         query_tokens = query_tokenization(query, self.tokenizer)
         query_embeds = query_embed_extraction(query_tokens, self.model, 
@@ -224,5 +255,6 @@ class SemanticSearch:
         pred = [self.index_keys[x] 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)
 
         return preds[:top], docs[:top], scores[:top]