Short query index

healthcheck_bm25.py

@@ -1,28 +1,32 @@
 from pathlib import Path
 
 from llama_index.retrievers.bm25 import BM25Retriever
-from main import extract_court_decision_text
 
 
-PERSIST_PATH = Path("Save_Index_Local")
-
-INDEX_NAME = "bm25_retriever"
-# INDEX_NAME = "bm25_retriever_meta"
 
-TEST_CD_URL = "https://reyestr.court.gov.ua/Review/118766467"
+# from main import extract_court_decision_text
+# INDEX_NAME = "bm25_retriever_long"
+# TEST_CD_URL = "https://reyestr.court.gov.ua/Review/118766467"
 # TEST_CD_URL = "https://reyestr.court.gov.ua/Review/118763429"
+# query = extract_court_decision_text(TEST_CD_URL)
 
-PRINT_CD = False
+INDEX_NAME = "bm25_retriever_short"
+query = (
+    "Викрадення майна, злочини, пов'язані з порушенням законодавчих норм щодо обігу та використання "
+    "документів, печаток, штампів, бланків, а також спеціальних технічних засобів "
+    "для отримання інформації та комунікаційних ліній."
+)
 
+PERSIST_PATH = Path("Save_Index_Local")
+PRINT_QUERY = True
 
-retriever = BM25Retriever.from_persist_dir(str(PERSIST_PATH / INDEX_NAME))
 
-court_decision_text = extract_court_decision_text(TEST_CD_URL)
+retriever = BM25Retriever.from_persist_dir(str(PERSIST_PATH / INDEX_NAME))
 
-if PRINT_CD:
-    print(court_decision_text, "\n\n\n\n\n")
+if PRINT_QUERY:
+    print(query)
 
-nodes_with_score = retriever.retrieve(court_decision_text)
+nodes_with_score = retriever.retrieve(query)
 for index, node_with_score in enumerate(nodes_with_score, start=1):
     source_title = node_with_score.node.metadata.get("title", "Невідомий заголовок")
     print(index, f"{node_with_score.score:.4f}", source_title, "\n", sep="\t")

init_bm25.py

@@ -10,15 +10,21 @@ from llama_index.core.vector_stores.utils import node_to_metadata_dict
 from llama_index.retrievers.bm25 import BM25Retriever
 
 
-
 PERSIST_PATH = Path("Save_Index_Local")
 LP_INFO_FILE = "legal_position_with_categories_documents_all.xlsx"
 
-INDEX_NAME = "bm25_retriever"
-USE_META = False
+# NOTE: previously named "bm25_retriever_meta"
+# INDEX_NAME = "bm25_retriever_long"
+# k1 = 1.88
+# b = 1.25
+# delta = 0.5
+# method = "robertson"
 
-# INDEX_NAME = "bm25_retriever_meta"
-# USE_META = True
+INDEX_NAME = "bm25_retriever_short"
+k1 = 0.35
+b = 0.6
+delta = 0.5
+method = "robertson"
 
 
 def clean_string(text: pd.Series):
@@ -69,49 +75,31 @@ def find_matching_pattern(categories):
 
 final_df = pd.read_excel(LP_INFO_FILE)
 
-if USE_META:
-    category_columns = [
-        col for col in final_df.columns if re.match(r"category_\d+$", col)
-    ]
+category_columns = [col for col in final_df.columns if re.match(r"category_\d+$", col)]
 
-    text_columns = ["title", "text_lp", "category_all"] + category_columns
-    final_df[text_columns] = final_df[text_columns].apply(clean_string)
+text_columns = ["title", "text_lp", "category_all"] + category_columns
+final_df[text_columns] = final_df[text_columns].apply(clean_string)
 
-    final_df["category_search"] = final_df[category_columns].apply(
-        lambda row: ", ".join([str(val) for val in row if pd.notna(val)]), axis=1
-    )
-    final_df["category_filter"] = final_df["category_all"].apply(find_matching_pattern)
+final_df["category_search"] = final_df[category_columns].apply(
+    lambda row: ", ".join([str(val) for val in row if pd.notna(val)]), axis=1
+)
+final_df["category_filter"] = final_df["category_all"].apply(find_matching_pattern)
 
-    legal_position_title_category = [
-        Document(
-            text=row["text_lp"],  # type: ignore
-            metadata={  # type: ignore
-                "lp_id": row["id"],
-                "title": row["title"],
-                "doc_id": row["document_ids"],
-                "category_filter": find_matching_pattern(row["category_all"]),
-                "category_search": row["category_search"],
-            },
-            excluded_embed_metadata_keys=["doc_id", "category_filter"],
-            excluded_llm_metadata_keys=["doc_id", "category_filter"],
-        )
-        for _, row in final_df.iterrows()
-    ]
-else:
-    final_df[["title", "text_lp"]] = final_df[["title", "text_lp"]].apply(clean_string)
-    legal_position_title_category = [
-        Document(
-            text=row["text_lp"],  # type: ignore
-            metadata={  # type: ignore
-                "lp_id": row["id"],
-                "doc_id": row["document_ids"],
-                "title": row["title"],
-            },
-            excluded_embed_metadata_keys=["lp_id", "doc_id", "title"],
-            excluded_llm_metadata_keys=["lp_id", "doc_id", "title"],
-        )
-        for _, row in final_df.iterrows()
-    ]
+legal_position_title_category = [
+    Document(
+        text=row["text_lp"],  # type: ignore
+        metadata={  # type: ignore
+            "lp_id": row["id"],
+            "title": row["title"],
+            "doc_id": row["document_ids"],
+            "category_filter": find_matching_pattern(row["category_all"]),
+            "category_search": row["category_search"],
+        },
+        excluded_embed_metadata_keys=["doc_id", "category_filter"],
+        excluded_llm_metadata_keys=["doc_id", "category_filter"],
+    )
+    for _, row in final_df.iterrows()
+]
 
 
 ukrainian_stopwords_1 = [
@@ -281,10 +269,10 @@ corpus_tokens = bm25s.tokenize(
 )
 
 existing_bm25 = bm25s.BM25(
-    k1=1.88,
-    b=1.25,
-    delta=0.5,
-    method="robertson",
+    k1=k1,
+    b=b,
+    delta=b,
+    method=method,
     # No corpus is saved without this line:
     corpus=corpus,  # stores metadata and prevents TypeError: 'NoneType' object is not subscriptable
 )