diff --git a/evaluation/evaluate_retrieval.py b/evaluation/evaluate_retrieval.py index 113f163..7847845 100644 --- a/evaluation/evaluate_retrieval.py +++ b/evaluation/evaluate_retrieval.py @@ -46,6 +46,11 @@ EVALUATION_MODES = ( "hybrid", ) +VALID_SPLITS = ( + "dev", + "test", +) + def load_questions( path: Path, @@ -75,6 +80,8 @@ def load_questions( dict[str, Any] ] = [] + seen_ids: set[str] = set() + for index, item in enumerate( data, start=1, @@ -96,6 +103,10 @@ def load_questions( "question" ) + split = item.get( + "split" + ) + expected_documents = item.get( "expected_documents" ) @@ -111,6 +122,15 @@ def load_questions( f"Položka {index} nemá platné id" ) + if question_id in seen_ids: + raise ValueError( + f"Dataset obsahuje duplicitné id: {question_id}" + ) + + seen_ids.add( + question_id + ) + if ( not isinstance( question, @@ -122,6 +142,12 @@ def load_questions( f"{question_id}: chýba otázka" ) + if split not in VALID_SPLITS: + raise ValueError( + f"{question_id}: split musí byť " + "'dev' alebo 'test'" + ) + if ( not isinstance( expected_documents, @@ -153,6 +179,45 @@ def load_questions( return questions +def filter_questions_by_split( + questions: list[ + dict[str, Any] + ], + split: str, +) -> list[dict[str, Any]]: + if split == "all": + return questions + + return [ + item + for item in questions + if item.get("split") == split + ] + + +def count_splits( + questions: list[ + dict[str, Any] + ], +) -> dict[str, int]: + counts = { + "dev": 0, + "test": 0, + } + + for item in questions: + split = item.get( + "split" + ) + + if split in counts: + counts[ + split + ] += 1 + + return counts + + def load_index_document_paths( db_file: Path, ) -> set[str]: @@ -497,11 +562,20 @@ def evaluate_question_mode( ) row: dict[str, Any] = { - "id": question["id"], + "id": question[ + "id" + ], + "split": question.get( + "split" + ), "category": question.get( "category", "unknown", ), + "difficulty": question.get( + "difficulty", + "unknown", + ), "question": question[ "question" ], @@ -515,7 +589,9 @@ def evaluate_question_mode( 6, ), "top_documents": [ - item["document_path"] + item[ + "document_path" + ] for item in ranked_documents ], "top_source_urls": [ @@ -695,11 +771,57 @@ def aggregate_by_category( } +def aggregate_by_difficulty( + rows: list[ + dict[str, Any] + ], +) -> dict[ + str, + dict[str, Any], +]: + grouped: dict[ + str, + list[dict[str, Any]], + ] = defaultdict( + list + ) + + for row in rows: + difficulty = str( + row.get( + "difficulty", + "unknown", + ) + ) + + grouped[ + difficulty + ].append( + row + ) + + return { + difficulty: aggregate_metrics( + difficulty_rows + ) + for ( + difficulty, + difficulty_rows, + ) in sorted( + grouped.items() + ) + } + + def print_summary( summary: dict[ str, dict[str, Any], ], + *, + split: str, + selected_count: int, + total_count: int, ) -> None: print() print( @@ -709,6 +831,19 @@ def print_summary( "=" * 78 ) + print( + f"Split: {split}" + ) + + print( + f"Questions: " + f"{selected_count}/{total_count}" + ) + + print( + "-" * 78 + ) + header = ( f"{'Mode':<10}" f"{'Questions':>10}" @@ -786,7 +921,9 @@ def save_csv_results( fieldnames = [ "id", + "split", "category", + "difficulty", "mode", "question", "first_relevant_rank", @@ -887,6 +1024,23 @@ def parse_args() -> argparse.Namespace: ), ) + parser.add_argument( + "--split", + choices=( + "dev", + "test", + "all", + ), + default="dev", + help=( + "Časť datasetu: " + "dev = ladenie, " + "test = finálne hodnotenie, " + "all = celý dataset. " + "Predvolené je dev." + ), + ) + parser.add_argument( "--limit", type=int, @@ -943,10 +1097,60 @@ def parse_args() -> argparse.Namespace: def main() -> None: args = parse_args() - questions = load_questions( + all_questions = load_questions( args.questions ) + split_counts = count_splits( + all_questions + ) + + questions = ( + filter_questions_by_split( + all_questions, + args.split, + ) + ) + + if not questions: + raise RuntimeError( + f"Pre split '{args.split}' " + "sa nenašli žiadne otázky." + ) + + print() + print( + "Dataset" + ) + print( + "=" * 60 + ) + + print( + f"Total: {len(all_questions)}" + ) + + print( + f"Dev: {split_counts['dev']}" + ) + + print( + f"Test: {split_counts['test']}" + ) + + print( + f"Selected split: {args.split}" + ) + + print( + f"Selected questions: {len(questions)}" + ) + + print( + "=" * 60 + ) + print() + indexed_documents = ( load_index_document_paths( args.db @@ -994,7 +1198,7 @@ def main() -> None: start=1, ): print( - f"[{index:02d}/{total:02d}] " + f"[{index:04d}/{total:04d}] " f"{question['id']}: " f"{question['question']}" ) @@ -1002,7 +1206,9 @@ def main() -> None: mode_results = ( retrieve_all_modes( args.db, - question["question"], + question[ + "question" + ], limit=args.limit, published_only=( args.published_only @@ -1037,11 +1243,21 @@ def main() -> None: ], ] = {} + by_difficulty: dict[ + str, + dict[ + str, + dict[str, Any], + ], + ] = {} + for mode in EVALUATION_MODES: mode_rows = [ row for row in evaluation_rows - if row["mode"] == mode + if row[ + "mode" + ] == mode ] summary[ @@ -1056,6 +1272,12 @@ def main() -> None: mode_rows ) + by_difficulty[ + mode + ] = aggregate_by_difficulty( + mode_rows + ) + payload = { "configuration": { "questions_file": str( @@ -1064,9 +1286,25 @@ def main() -> None: "database": str( args.db ), - "question_count": len( + "split": ( + args.split + ), + "dataset_question_count": len( + all_questions + ), + "selected_question_count": len( questions ), + "dev_question_count": ( + split_counts[ + "dev" + ] + ), + "test_question_count": ( + split_counts[ + "test" + ] + ), "limit": args.limit, "published_only": ( args.published_only @@ -1097,18 +1335,35 @@ def main() -> None: ), }, "summary": summary, - "by_category": by_category, - "questions": evaluation_rows, + "by_category": ( + by_category + ), + "by_difficulty": ( + by_difficulty + ), + "questions": ( + evaluation_rows + ), } + filename_suffix = ( + args.split + ) + json_path = ( args.output_dir - / "retrieval_results.json" + / ( + "retrieval_results_" + f"{filename_suffix}.json" + ) ) csv_path = ( args.output_dir - / "retrieval_results.csv" + / ( + "retrieval_results_" + f"{filename_suffix}.csv" + ) ) save_json_results( @@ -1122,15 +1377,24 @@ def main() -> None: ) print_summary( - summary + summary, + split=args.split, + selected_count=len( + questions + ), + total_count=len( + all_questions + ), ) print( "Výsledky:" ) + print( f" JSON: {json_path}" ) + print( f" CSV: {csv_path}" ) diff --git a/scripts/search_core.py b/scripts/search_core.py new file mode 100644 index 0000000..fda8e5f --- /dev/null +++ b/scripts/search_core.py @@ -0,0 +1,1179 @@ +from __future__ import annotations + +import json +import re +import sqlite3 +import unicodedata +from collections import defaultdict +from typing import Any + +from scripts.embedding_utils import ( + blob_to_vector, + cosine_similarity, + embed_query, +) + + +WORD_RE = re.compile( + r"[^\W_]+", + re.UNICODE, +) + + +BM25_WEIGHTS = ( + 0.0, + 10.0, + 7.0, + 5.0, + 8.0, + 4.0, + 1.0, +) + +BM25_SQL = ", ".join( + str(value) + for value in BM25_WEIGHTS +) + + +DEFAULT_CANDIDATE_MULTIPLIER = 8 +MIN_CANDIDATES = 50 +MIN_STEM_PREFIX_LENGTH = 5 + +RRF_K = 60 +FTS_RRF_WEIGHT = 1.0 +VECTOR_RRF_WEIGHT = 1.5 +ANY_TERM_RRF_WEIGHT = 0.25 + + +STRATEGY_PRIORITY = { + "all_terms": 3, + "prefix_terms": 2, + "any_term": 1, +} + + +def normalize_for_compare( + text: str, +) -> str: + text = unicodedata.normalize( + "NFKD", + text.casefold(), + ) + + text = "".join( + character + for character in text + if not unicodedata.combining( + character + ) + ) + + return " ".join( + WORD_RE.findall(text) + ) + + +def query_tokens( + query: str, +) -> list[str]: + tokens: list[str] = [] + seen: set[str] = set() + + for token in WORD_RE.findall( + query + ): + normalized = normalize_for_compare( + token + ) + + if not normalized: + continue + + if normalized in seen: + continue + + tokens.append( + token + ) + + seen.add( + normalized + ) + + return tokens + + +def quote_fts_token( + token: str, + *, + use_prefix: bool = True, + shorten: bool = False, +) -> str: + value = token + + if ( + shorten + and len(value) + > MIN_STEM_PREFIX_LENGTH + ): + value = value[ + :MIN_STEM_PREFIX_LENGTH + ] + + escaped = value.replace( + '"', + '""', + ) + + suffix = ( + "*" + if ( + use_prefix + and len(value) >= 4 + ) + else "" + ) + + return f'"{escaped}"{suffix}' + + +def build_match_queries( + query: str, +) -> list[tuple[str, str]]: + tokens = query_tokens( + query + ) + + if not tokens: + return [] + + full_terms = [ + quote_fts_token( + token + ) + for token in tokens + ] + + all_terms_query = ( + " AND ".join( + full_terms + ) + ) + + queries = [ + ( + "all_terms", + all_terms_query, + ) + ] + + shortened_terms = [ + quote_fts_token( + token, + shorten=True, + ) + for token in tokens + ] + + shortened_query = ( + " AND ".join( + shortened_terms + ) + ) + + if ( + shortened_query + != all_terms_query + ): + queries.append( + ( + "prefix_terms", + shortened_query, + ) + ) + + if len(full_terms) > 1: + queries.append( + ( + "any_term", + " OR ".join( + full_terms + ), + ) + ) + + return queries + + +def verify_search_schema( + conn: sqlite3.Connection, +) -> None: + fts_row = conn.execute( + """ + SELECT 1 + FROM sqlite_master + WHERE type = 'table' + AND name = 'chunks_fts' + """ + ).fetchone() + + if fts_row is None: + raise RuntimeError( + "FTS5 index v databáze chýba. " + "Spusti python scripts/rebuild_index.py." + ) + + embedding_row = conn.execute( + """ + SELECT 1 + FROM sqlite_master + WHERE type = 'table' + AND name = 'chunk_embeddings' + """ + ).fetchone() + + if embedding_row is None: + raise RuntimeError( + "Embedding index v databáze chýba. " + "Spusti python scripts/rebuild_index.py." + ) + + +def embedding_index_info( + conn: sqlite3.Connection, +) -> tuple[str, int]: + row = conn.execute( + """ + SELECT + model, + dimensions + FROM chunk_embeddings + LIMIT 1 + """ + ).fetchone() + + if row is None: + raise RuntimeError( + "Embedding index je prázdny" + ) + + return ( + str(row["model"]), + int(row["dimensions"]), + ) + + +def make_source_url( + document_path: str, +) -> str: + clean_path = document_path + + if clean_path.startswith( + "pages/" + ): + clean_path = clean_path[ + len("pages/"): + ] + + if clean_path.endswith( + "/README.md" + ): + clean_path = clean_path[ + :-len("/README.md") + ] + + return ( + "https://zp.kemt.fei.tuke.sk/" + f"{clean_path}" + ) + + +def load_labels( + conn: sqlite3.Connection, + table: str, + column: str, + chunk_ids: list[str], +) -> dict[str, list[str]]: + if not chunk_ids: + return {} + + placeholders = ",".join( + "?" + for _ in chunk_ids + ) + + rows = conn.execute( + f""" + SELECT + chunk_id, + {column} + FROM {table} + WHERE chunk_id IN ( + {placeholders} + ) + ORDER BY + chunk_id, + {column} + """, + chunk_ids, + ).fetchall() + + values: dict[ + str, + list[str], + ] = defaultdict( + list + ) + + for chunk_id, value in rows: + values[ + chunk_id + ].append( + value + ) + + return dict( + values + ) + + +def run_fts_query( + conn: sqlite3.Connection, + match_query: str, + candidate_limit: int, + published_only: bool, +) -> list[dict[str, Any]]: + rows = conn.execute( + f""" + SELECT + chunks.chunk_id, + chunks.document_path, + chunks.title, + chunks.author, + chunks.published, + chunks.chunk_index, + chunks.heading_paths_json, + chunks.text, + chunks.text_length, + chunks.token_count, + chunks.content_hash, + chunks_fts.rank AS bm25_score, + snippet( + chunks_fts, + 6, + '', + '', + ' … ', + 36 + ) AS snippet + FROM chunks_fts + JOIN chunks + ON chunks.id = chunks_fts.rowid + WHERE chunks_fts MATCH ? + AND chunks_fts.rank MATCH + 'bm25({BM25_SQL})' + AND ( + ? = 0 + OR chunks.published = 1 + ) + ORDER BY + chunks_fts.rank ASC, + chunks.id ASC + LIMIT ? + """, + ( + match_query, + 1 if published_only else 0, + candidate_limit, + ), + ).fetchall() + + return [ + dict(row) + for row in rows + ] + + +def run_vector_query( + conn: sqlite3.Connection, + query: str, + candidate_limit: int, + published_only: bool, +) -> list[dict[str, Any]]: + model_name, dimensions = ( + embedding_index_info( + conn + ) + ) + + query_vector = embed_query( + query, + model_name=model_name, + ) + + if ( + int(query_vector.shape[0]) + != dimensions + ): + raise RuntimeError( + "Rozmer query embeddingu " + "sa nezhoduje s indexom" + ) + + rows = conn.execute( + """ + SELECT + chunks.chunk_id, + chunks.document_path, + chunks.title, + chunks.author, + chunks.published, + chunks.chunk_index, + chunks.heading_paths_json, + chunks.text, + chunks.text_length, + chunks.token_count, + chunks.content_hash, + chunk_embeddings.embedding, + chunk_embeddings.dimensions + FROM chunk_embeddings + JOIN chunks + ON chunks.chunk_id + = chunk_embeddings.chunk_id + WHERE ( + ? = 0 + OR chunks.published = 1 + ) + """, + ( + 1 if published_only else 0, + ), + ).fetchall() + + candidates: list[ + dict[str, Any] + ] = [] + + for row in rows: + item = dict( + row + ) + + blob = item.pop( + "embedding" + ) + + stored_dimensions = int( + item.pop( + "dimensions" + ) + ) + + vector = blob_to_vector( + blob, + stored_dimensions, + ) + + similarity = cosine_similarity( + query_vector, + vector, + ) + + item["vector_score"] = round( + similarity, + 6, + ) + + candidates.append( + item + ) + + candidates.sort( + key=lambda item: ( + -item["vector_score"], + item["document_path"], + item["chunk_index"], + ) + ) + + return candidates[ + :candidate_limit + ] + + +def tokens_match_for_title( + title_token: str, + query_token: str, +) -> bool: + if ( + not title_token + or not query_token + ): + return False + + if title_token == query_token: + return True + + shorter_length = min( + len(title_token), + len(query_token), + ) + + # Pri veľmi krátkych tokenoch nechceme + # agresívne prefixové zhody. + if shorter_length < 3: + return False + + return ( + title_token.startswith( + query_token + ) + or query_token.startswith( + title_token + ) + ) + + +def query_contains_title( + query: str, + title: str, +) -> bool: + normalized_query = normalize_for_compare( + query + ) + + normalized_title = normalize_for_compare( + title + ) + + if ( + not normalized_query + or not normalized_title + ): + return False + + title_tokens = ( + normalized_title.split() + ) + + query_token_values = ( + normalized_query.split() + ) + + # Jednoslovné názvy by mohli vytvárať + # príliš veľa falošných boostov. + if len(title_tokens) < 2: + return False + + return all( + any( + tokens_match_for_title( + title_token, + query_token, + ) + for query_token + in query_token_values + ) + for title_token + in title_tokens + ) + + +def exact_match_bonus( + query: str, + item: dict[str, Any], + tags: list[str], + categories: list[str], +) -> float: + normalized_query = normalize_for_compare( + query + ) + + if not normalized_query: + return 0.0 + + raw_title = ( + item.get("title") + or "" + ) + + title = normalize_for_compare( + raw_title + ) + + author = normalize_for_compare( + item.get("author") + or "" + ) + + path = normalize_for_compare( + item.get("document_path") + or "" + ) + + text = normalize_for_compare( + item.get("text") + or "" + ) + + normalized_tags = [ + normalize_for_compare( + value + ) + for value in tags + ] + + normalized_categories = [ + normalize_for_compare( + value + ) + for value in categories + ] + + bonus = 0.0 + + if title == normalized_query: + bonus += 6.0 + + elif query_contains_title( + query, + raw_title, + ): + bonus += 5.0 + + elif normalized_query in title: + bonus += 3.0 + + if author == normalized_query: + bonus += 5.0 + + elif normalized_query in author: + bonus += 2.0 + + if normalized_query in path: + bonus += 2.0 + + if normalized_query in ( + normalized_tags + ): + bonus += 4.0 + + if normalized_query in ( + normalized_categories + ): + bonus += 3.0 + + if normalized_query in text: + bonus += 1.5 + + return bonus + + +def database_bool( + value: Any, +) -> bool | None: + if value is None: + return None + + return bool( + value + ) + + +def parse_heading_paths( + item: dict[str, Any], +) -> list: + try: + return json.loads( + item.pop( + "heading_paths_json" + ) + or "[]" + ) + + except json.JSONDecodeError: + return [] + + +def load_result_labels( + conn: sqlite3.Connection, + candidates: list[ + dict[str, Any] + ], +) -> tuple[ + dict[str, list[str]], + dict[str, list[str]], +]: + chunk_ids = [ + item["chunk_id"] + for item in candidates + ] + + tags = load_labels( + conn, + "chunk_tags", + "tag", + chunk_ids, + ) + + categories = load_labels( + conn, + "chunk_categories", + "category", + chunk_ids, + ) + + return ( + tags, + categories, + ) + + +def add_fts_metadata( + conn: sqlite3.Connection, + query: str, + candidates: list[ + dict[str, Any] + ], +) -> list[dict[str, Any]]: + ( + tags_by_chunk, + categories_by_chunk, + ) = load_result_labels( + conn, + candidates, + ) + + results: list[ + dict[str, Any] + ] = [] + + for item in candidates: + chunk_id = item[ + "chunk_id" + ] + + tags = tags_by_chunk.get( + chunk_id, + [], + ) + + categories = ( + categories_by_chunk.get( + chunk_id, + [], + ) + ) + + bm25_score = float( + item.pop( + "bm25_score" + ) + ) + + strategy = item.pop( + "strategy" + ) + + base_score = max( + 0.0, + -bm25_score, + ) + + score = ( + base_score + + exact_match_bonus( + query, + item, + tags, + categories, + ) + ) + + heading_paths = ( + parse_heading_paths( + item + ) + ) + + item["published"] = ( + database_bool( + item.get( + "published" + ) + ) + ) + + item["_strategy_priority"] = ( + STRATEGY_PRIORITY[ + strategy + ] + ) + + item.update( + { + "heading_paths": ( + heading_paths + ), + "tags": tags, + "categories": ( + categories + ), + "score": round( + score, + 6, + ), + "bm25_score": round( + bm25_score, + 6, + ), + "match_strategy": ( + strategy + ), + "source_url": ( + make_source_url( + item[ + "document_path" + ] + ) + ), + } + ) + + results.append( + item + ) + + results.sort( + key=lambda item: ( + -item[ + "_strategy_priority" + ], + -item["score"], + item["bm25_score"], + item["document_path"], + item["chunk_index"], + ) + ) + + for item in results: + item.pop( + "_strategy_priority", + None, + ) + + return results + + +def add_vector_metadata( + conn: sqlite3.Connection, + candidates: list[ + dict[str, Any] + ], +) -> list[dict[str, Any]]: + ( + tags_by_chunk, + categories_by_chunk, + ) = load_result_labels( + conn, + candidates, + ) + + results: list[ + dict[str, Any] + ] = [] + + for item in candidates: + chunk_id = item[ + "chunk_id" + ] + + item["published"] = ( + database_bool( + item.get( + "published" + ) + ) + ) + + item["heading_paths"] = ( + parse_heading_paths( + item + ) + ) + + item["tags"] = ( + tags_by_chunk.get( + chunk_id, + [], + ) + ) + + item["categories"] = ( + categories_by_chunk.get( + chunk_id, + [], + ) + ) + + item["source_url"] = ( + make_source_url( + item[ + "document_path" + ] + ) + ) + + text = ( + item.get("text") + or "" + ) + + item["snippet"] = ( + text[ + :320 + ].strip() + ) + + results.append( + item + ) + + return results + + +def fuse_hybrid_results( + fts_results: list[ + dict[str, Any] + ], + vector_results: list[ + dict[str, Any] + ], +) -> list[dict[str, Any]]: + merged: dict[ + str, + dict[str, Any], + ] = {} + + scores: dict[ + str, + float, + ] = defaultdict( + float + ) + + fts_ranks: dict[ + str, + int, + ] = {} + + vector_ranks: dict[ + str, + int, + ] = {} + + for rank, item in enumerate( + fts_results, + start=1, + ): + chunk_id = item[ + "chunk_id" + ] + + merged[ + chunk_id + ] = dict( + item + ) + + fts_ranks[ + chunk_id + ] = rank + + strategy = item.get( + "match_strategy" + ) + + fts_weight = ( + ANY_TERM_RRF_WEIGHT + if strategy + == "any_term" + else FTS_RRF_WEIGHT + ) + + scores[ + chunk_id + ] += ( + fts_weight + / ( + RRF_K + + rank + ) + ) + + for rank, item in enumerate( + vector_results, + start=1, + ): + chunk_id = item[ + "chunk_id" + ] + + vector_ranks[ + chunk_id + ] = rank + + scores[ + chunk_id + ] += ( + VECTOR_RRF_WEIGHT + / ( + RRF_K + + rank + ) + ) + + if chunk_id not in merged: + merged[ + chunk_id + ] = dict( + item + ) + + else: + merged[ + chunk_id + ][ + "vector_score" + ] = item[ + "vector_score" + ] + + results: list[ + dict[str, Any] + ] = [] + + for chunk_id, item in ( + merged.items() + ): + item["fts_score"] = ( + item.get( + "score" + ) + ) + + item["fts_rank"] = ( + fts_ranks.get( + chunk_id + ) + ) + + item["vector_rank"] = ( + vector_ranks.get( + chunk_id + ) + ) + + item.setdefault( + "vector_score", + None, + ) + + item.setdefault( + "bm25_score", + None, + ) + + item.setdefault( + "match_strategy", + None, + ) + + hybrid_score = ( + scores[ + chunk_id + ] + ) + + item["hybrid_score"] = ( + round( + hybrid_score, + 8, + ) + ) + + item["score"] = round( + hybrid_score, + 8, + ) + + results.append( + item + ) + + results.sort( + key=lambda item: ( + -item[ + "hybrid_score" + ], + item[ + "document_path" + ], + item[ + "chunk_index" + ], + ) + ) + + return results + + +def diversify_results( + results: list[ + dict[str, Any] + ], + limit: int, + max_per_document: int, +) -> list[dict[str, Any]]: + if max_per_document <= 0: + return results[ + :limit + ] + + selected: list[ + dict[str, Any] + ] = [] + + document_counts: dict[ + str, + int, + ] = defaultdict( + int + ) + + for item in results: + document_path = item[ + "document_path" + ] + + if ( + document_counts[ + document_path + ] + >= max_per_document + ): + continue + + selected.append( + item + ) + + document_counts[ + document_path + ] += 1 + + if len(selected) >= limit: + break + + return selected diff --git a/scripts/search_utils.py b/scripts/search_utils.py index fd7f045..3198bd3 100644 --- a/scripts/search_utils.py +++ b/scripts/search_utils.py @@ -1,1059 +1,48 @@ from __future__ import annotations -import json -import re import sqlite3 -import unicodedata -from collections import defaultdict from pathlib import Path from typing import Any -from scripts.embedding_utils import ( - blob_to_vector, - cosine_similarity, - embed_query, +# Tieto importy sú zámerne verejné. +# Zachovávajú spätnú kompatibilitu pre testy +# a evaluačné skripty, ktoré ich importujú +# zo scripts.search_utils. +from scripts.search_core import ( + ANY_TERM_RRF_WEIGHT, + BM25_SQL, + BM25_WEIGHTS, + DEFAULT_CANDIDATE_MULTIPLIER, + FTS_RRF_WEIGHT, + MIN_CANDIDATES, + MIN_STEM_PREFIX_LENGTH, + RRF_K, + STRATEGY_PRIORITY, + VECTOR_RRF_WEIGHT, + WORD_RE, + add_fts_metadata, + add_vector_metadata, + build_match_queries, + database_bool, + diversify_results, + embedding_index_info, + exact_match_bonus, + fuse_hybrid_results, + load_labels, + load_result_labels, + make_source_url, + normalize_for_compare, + parse_heading_paths, + query_contains_title, + query_tokens, + quote_fts_token, + run_fts_query, + run_vector_query, + tokens_match_for_title, + verify_search_schema, ) -WORD_RE = re.compile( - r"[^\W_]+", - re.UNICODE, -) - - -BM25_WEIGHTS = ( - 0.0, - 10.0, - 7.0, - 5.0, - 8.0, - 4.0, - 1.0, -) - -BM25_SQL = ", ".join( - str(value) - for value in BM25_WEIGHTS -) - - -DEFAULT_CANDIDATE_MULTIPLIER = 8 -MIN_CANDIDATES = 50 -MIN_STEM_PREFIX_LENGTH = 5 - -RRF_K = 60 -FTS_RRF_WEIGHT = 1.0 -VECTOR_RRF_WEIGHT = 1.5 -ANY_TERM_RRF_WEIGHT = 0.25 - - -STRATEGY_PRIORITY = { - "all_terms": 3, - "prefix_terms": 2, - "any_term": 1, -} - - -def normalize_for_compare( - text: str, -) -> str: - text = unicodedata.normalize( - "NFKD", - text.casefold(), - ) - - text = "".join( - character - for character in text - if not unicodedata.combining( - character - ) - ) - - return " ".join( - WORD_RE.findall(text) - ) - - -def query_tokens( - query: str, -) -> list[str]: - tokens: list[str] = [] - seen: set[str] = set() - - for token in WORD_RE.findall( - query - ): - normalized = normalize_for_compare( - token - ) - - if not normalized: - continue - - if normalized in seen: - continue - - tokens.append( - token - ) - - seen.add( - normalized - ) - - return tokens - - -def quote_fts_token( - token: str, - *, - use_prefix: bool = True, - shorten: bool = False, -) -> str: - value = token - - if ( - shorten - and len(value) - > MIN_STEM_PREFIX_LENGTH - ): - value = value[ - :MIN_STEM_PREFIX_LENGTH - ] - - escaped = value.replace( - '"', - '""', - ) - - suffix = ( - "*" - if ( - use_prefix - and len(value) >= 4 - ) - else "" - ) - - return f'"{escaped}"{suffix}' - - -def build_match_queries( - query: str, -) -> list[tuple[str, str]]: - tokens = query_tokens( - query - ) - - if not tokens: - return [] - - full_terms = [ - quote_fts_token( - token - ) - for token in tokens - ] - - all_terms_query = ( - " AND ".join( - full_terms - ) - ) - - queries = [ - ( - "all_terms", - all_terms_query, - ) - ] - - shortened_terms = [ - quote_fts_token( - token, - shorten=True, - ) - for token in tokens - ] - - shortened_query = ( - " AND ".join( - shortened_terms - ) - ) - - if ( - shortened_query - != all_terms_query - ): - queries.append( - ( - "prefix_terms", - shortened_query, - ) - ) - - if len(full_terms) > 1: - queries.append( - ( - "any_term", - " OR ".join( - full_terms - ), - ) - ) - - return queries - - -def verify_search_schema( - conn: sqlite3.Connection, -) -> None: - fts_row = conn.execute( - """ - SELECT 1 - FROM sqlite_master - WHERE type = 'table' - AND name = 'chunks_fts' - """ - ).fetchone() - - if fts_row is None: - raise RuntimeError( - "FTS5 index v databáze chýba. " - "Spusti python scripts/rebuild_index.py." - ) - - embedding_row = conn.execute( - """ - SELECT 1 - FROM sqlite_master - WHERE type = 'table' - AND name = 'chunk_embeddings' - """ - ).fetchone() - - if embedding_row is None: - raise RuntimeError( - "Embedding index v databáze chýba. " - "Spusti python scripts/rebuild_index.py." - ) - - -def embedding_index_info( - conn: sqlite3.Connection, -) -> tuple[str, int]: - row = conn.execute( - """ - SELECT - model, - dimensions - FROM chunk_embeddings - LIMIT 1 - """ - ).fetchone() - - if row is None: - raise RuntimeError( - "Embedding index je prázdny" - ) - - return ( - str(row["model"]), - int(row["dimensions"]), - ) - - -def make_source_url( - document_path: str, -) -> str: - clean_path = document_path - - if clean_path.startswith( - "pages/" - ): - clean_path = clean_path[ - len("pages/"): - ] - - if clean_path.endswith( - "/README.md" - ): - clean_path = clean_path[ - :-len("/README.md") - ] - - return ( - "https://zp.kemt.fei.tuke.sk/" - f"{clean_path}" - ) - - -def load_labels( - conn: sqlite3.Connection, - table: str, - column: str, - chunk_ids: list[str], -) -> dict[str, list[str]]: - if not chunk_ids: - return {} - - placeholders = ",".join( - "?" - for _ in chunk_ids - ) - - rows = conn.execute( - f""" - SELECT - chunk_id, - {column} - FROM {table} - WHERE chunk_id IN ( - {placeholders} - ) - ORDER BY - chunk_id, - {column} - """, - chunk_ids, - ).fetchall() - - values: dict[ - str, - list[str], - ] = defaultdict( - list - ) - - for chunk_id, value in rows: - values[ - chunk_id - ].append( - value - ) - - return dict( - values - ) - - -def run_fts_query( - conn: sqlite3.Connection, - match_query: str, - candidate_limit: int, - published_only: bool, -) -> list[dict[str, Any]]: - rows = conn.execute( - f""" - SELECT - chunks.chunk_id, - chunks.document_path, - chunks.title, - chunks.author, - chunks.published, - chunks.chunk_index, - chunks.heading_paths_json, - chunks.text, - chunks.text_length, - chunks.token_count, - chunks.content_hash, - chunks_fts.rank AS bm25_score, - snippet( - chunks_fts, - 6, - '', - '', - ' … ', - 36 - ) AS snippet - FROM chunks_fts - JOIN chunks - ON chunks.id = chunks_fts.rowid - WHERE chunks_fts MATCH ? - AND chunks_fts.rank MATCH - 'bm25({BM25_SQL})' - AND ( - ? = 0 - OR chunks.published = 1 - ) - ORDER BY - chunks_fts.rank ASC, - chunks.id ASC - LIMIT ? - """, - ( - match_query, - 1 if published_only else 0, - candidate_limit, - ), - ).fetchall() - - return [ - dict(row) - for row in rows - ] - - -def run_vector_query( - conn: sqlite3.Connection, - query: str, - candidate_limit: int, - published_only: bool, -) -> list[dict[str, Any]]: - model_name, dimensions = ( - embedding_index_info( - conn - ) - ) - - query_vector = embed_query( - query, - model_name=model_name, - ) - - if ( - int(query_vector.shape[0]) - != dimensions - ): - raise RuntimeError( - "Rozmer query embeddingu " - "sa nezhoduje s indexom" - ) - - rows = conn.execute( - """ - SELECT - chunks.chunk_id, - chunks.document_path, - chunks.title, - chunks.author, - chunks.published, - chunks.chunk_index, - chunks.heading_paths_json, - chunks.text, - chunks.text_length, - chunks.token_count, - chunks.content_hash, - chunk_embeddings.embedding, - chunk_embeddings.dimensions - FROM chunk_embeddings - JOIN chunks - ON chunks.chunk_id - = chunk_embeddings.chunk_id - WHERE ( - ? = 0 - OR chunks.published = 1 - ) - """, - ( - 1 if published_only else 0, - ), - ).fetchall() - - candidates: list[ - dict[str, Any] - ] = [] - - for row in rows: - item = dict( - row - ) - - blob = item.pop( - "embedding" - ) - - stored_dimensions = int( - item.pop( - "dimensions" - ) - ) - - vector = blob_to_vector( - blob, - stored_dimensions, - ) - - similarity = cosine_similarity( - query_vector, - vector, - ) - - item["vector_score"] = round( - similarity, - 6, - ) - - candidates.append( - item - ) - - candidates.sort( - key=lambda item: ( - -item["vector_score"], - item["document_path"], - item["chunk_index"], - ) - ) - - return candidates[ - :candidate_limit - ] - - -def exact_match_bonus( - query: str, - item: dict[str, Any], - tags: list[str], - categories: list[str], -) -> float: - normalized_query = normalize_for_compare( - query - ) - - if not normalized_query: - return 0.0 - - title = normalize_for_compare( - item.get("title") or "" - ) - - author = normalize_for_compare( - item.get("author") or "" - ) - - path = normalize_for_compare( - item.get("document_path") - or "" - ) - - text = normalize_for_compare( - item.get("text") or "" - ) - - normalized_tags = [ - normalize_for_compare( - value - ) - for value in tags - ] - - normalized_categories = [ - normalize_for_compare( - value - ) - for value in categories - ] - - bonus = 0.0 - - if title == normalized_query: - bonus += 6.0 - - elif normalized_query in title: - bonus += 3.0 - - if author == normalized_query: - bonus += 5.0 - - elif normalized_query in author: - bonus += 2.0 - - if normalized_query in path: - bonus += 2.0 - - if normalized_query in ( - normalized_tags - ): - bonus += 4.0 - - if normalized_query in ( - normalized_categories - ): - bonus += 3.0 - - if normalized_query in text: - bonus += 1.5 - - return bonus - - -def database_bool( - value: Any, -) -> bool | None: - if value is None: - return None - - return bool( - value - ) - - -def parse_heading_paths( - item: dict[str, Any], -) -> list: - try: - return json.loads( - item.pop( - "heading_paths_json" - ) - or "[]" - ) - - except json.JSONDecodeError: - return [] - - -def load_result_labels( - conn: sqlite3.Connection, - candidates: list[ - dict[str, Any] - ], -) -> tuple[ - dict[str, list[str]], - dict[str, list[str]], -]: - chunk_ids = [ - item["chunk_id"] - for item in candidates - ] - - tags = load_labels( - conn, - "chunk_tags", - "tag", - chunk_ids, - ) - - categories = load_labels( - conn, - "chunk_categories", - "category", - chunk_ids, - ) - - return ( - tags, - categories, - ) - - -def add_fts_metadata( - conn: sqlite3.Connection, - query: str, - candidates: list[ - dict[str, Any] - ], -) -> list[dict[str, Any]]: - ( - tags_by_chunk, - categories_by_chunk, - ) = load_result_labels( - conn, - candidates, - ) - - results: list[ - dict[str, Any] - ] = [] - - for item in candidates: - chunk_id = item[ - "chunk_id" - ] - - tags = tags_by_chunk.get( - chunk_id, - [], - ) - - categories = ( - categories_by_chunk.get( - chunk_id, - [], - ) - ) - - bm25_score = float( - item.pop( - "bm25_score" - ) - ) - - strategy = item.pop( - "strategy" - ) - - base_score = max( - 0.0, - -bm25_score, - ) - - score = ( - base_score - + exact_match_bonus( - query, - item, - tags, - categories, - ) - ) - - heading_paths = ( - parse_heading_paths( - item - ) - ) - - item["published"] = database_bool( - item.get( - "published" - ) - ) - - item["_strategy_priority"] = ( - STRATEGY_PRIORITY[ - strategy - ] - ) - - item.update( - { - "heading_paths": heading_paths, - "tags": tags, - "categories": categories, - "score": round( - score, - 6, - ), - "bm25_score": round( - bm25_score, - 6, - ), - "match_strategy": strategy, - "source_url": make_source_url( - item["document_path"] - ), - } - ) - - results.append( - item - ) - - results.sort( - key=lambda item: ( - -item["_strategy_priority"], - -item["score"], - item["bm25_score"], - item["document_path"], - item["chunk_index"], - ) - ) - - for item in results: - item.pop( - "_strategy_priority", - None, - ) - - return results - - -def add_vector_metadata( - conn: sqlite3.Connection, - candidates: list[ - dict[str, Any] - ], -) -> list[dict[str, Any]]: - ( - tags_by_chunk, - categories_by_chunk, - ) = load_result_labels( - conn, - candidates, - ) - - results: list[ - dict[str, Any] - ] = [] - - for item in candidates: - chunk_id = item[ - "chunk_id" - ] - - item["published"] = database_bool( - item.get( - "published" - ) - ) - - item["heading_paths"] = ( - parse_heading_paths( - item - ) - ) - - item["tags"] = ( - tags_by_chunk.get( - chunk_id, - [], - ) - ) - - item["categories"] = ( - categories_by_chunk.get( - chunk_id, - [], - ) - ) - - item["source_url"] = make_source_url( - item["document_path"] - ) - - text = ( - item.get("text") - or "" - ) - - item["snippet"] = ( - text[:320].strip() - ) - - results.append( - item - ) - - return results - - -def fuse_hybrid_results( - fts_results: list[ - dict[str, Any] - ], - vector_results: list[ - dict[str, Any] - ], -) -> list[dict[str, Any]]: - merged: dict[ - str, - dict[str, Any], - ] = {} - - scores: dict[ - str, - float, - ] = defaultdict( - float - ) - - fts_ranks: dict[ - str, - int, - ] = {} - - vector_ranks: dict[ - str, - int, - ] = {} - - for rank, item in enumerate( - fts_results, - start=1, - ): - chunk_id = item[ - "chunk_id" - ] - - merged[ - chunk_id - ] = dict( - item - ) - - fts_ranks[ - chunk_id - ] = rank - - strategy = item.get( - "match_strategy" - ) - - fts_weight = ( - ANY_TERM_RRF_WEIGHT - if strategy == "any_term" - else FTS_RRF_WEIGHT - ) - - scores[ - chunk_id - ] += ( - fts_weight - / ( - RRF_K - + rank - ) - ) - - for rank, item in enumerate( - vector_results, - start=1, - ): - chunk_id = item[ - "chunk_id" - ] - - vector_ranks[ - chunk_id - ] = rank - - scores[ - chunk_id - ] += ( - VECTOR_RRF_WEIGHT - / ( - RRF_K - + rank - ) - ) - - if chunk_id not in merged: - merged[ - chunk_id - ] = dict( - item - ) - - else: - merged[ - chunk_id - ][ - "vector_score" - ] = item[ - "vector_score" - ] - - results: list[ - dict[str, Any] - ] = [] - - for chunk_id, item in ( - merged.items() - ): - item["fts_score"] = ( - item.get("score") - ) - - item["fts_rank"] = ( - fts_ranks.get( - chunk_id - ) - ) - - item["vector_rank"] = ( - vector_ranks.get( - chunk_id - ) - ) - - item.setdefault( - "vector_score", - None, - ) - - item.setdefault( - "bm25_score", - None, - ) - - item.setdefault( - "match_strategy", - None, - ) - - hybrid_score = scores[ - chunk_id - ] - - item["hybrid_score"] = round( - hybrid_score, - 8, - ) - - item["score"] = round( - hybrid_score, - 8, - ) - - results.append( - item - ) - - results.sort( - key=lambda item: ( - -item["hybrid_score"], - item["document_path"], - item["chunk_index"], - ) - ) - - return results - - -def diversify_results( - results: list[ - dict[str, Any] - ], - limit: int, - max_per_document: int, -) -> list[dict[str, Any]]: - if max_per_document <= 0: - return results[ - :limit - ] - - selected: list[ - dict[str, Any] - ] = [] - - document_counts: dict[ - str, - int, - ] = defaultdict( - int - ) - - for item in results: - document_path = item[ - "document_path" - ] - - if ( - document_counts[ - document_path - ] - >= max_per_document - ): - continue - - selected.append( - item - ) - - document_counts[ - document_path - ] += 1 - - if len(selected) >= limit: - break - - return selected - - def search_database( db_file: Path, query: str, @@ -1078,8 +67,10 @@ def search_database( "results": [], } - match_queries = build_match_queries( - clean_query + match_queries = ( + build_match_queries( + clean_query + ) ) candidate_limit = max( @@ -1106,6 +97,10 @@ def search_database( conn ) + # ------------------------- + # FTS5 kandidáti + # ------------------------- + fts_candidates: list[ dict[str, Any] ] = [] @@ -1139,27 +134,50 @@ def search_database( strategy ] + # Používa sa prvá stratégia, + # ktorá vráti výsledky. break - fts_results = add_fts_metadata( - conn, - clean_query, - fts_candidates, + fts_results = ( + add_fts_metadata( + conn, + clean_query, + fts_candidates, + ) ) - vector_candidates = run_vector_query( - conn, - clean_query, - candidate_limit, - published_only, + # ------------------------- + # Embedding kandidáti + # ------------------------- + + vector_candidates = ( + run_vector_query( + conn, + clean_query, + candidate_limit, + published_only, + ) ) - vector_results = add_vector_metadata( - conn, - vector_candidates, + vector_results = ( + add_vector_metadata( + conn, + vector_candidates, + ) ) - # Presné FTS výsledky. + # ------------------------- + # Presné FTS výsledky + # ------------------------- + # + # Pri all_terms alebo prefix_terms + # embeddingy iba preradia chunky, + # ktoré už našiel FTS. + # + # Tým sa zabráni pridávaniu + # sémanticky podobných, ale + # lexikálne nesúvisiacich výsledkov. + if ( used_strategies and used_strategies[0] @@ -1169,22 +187,46 @@ def search_database( } ): fts_chunk_ids = { - item["chunk_id"] - for item in fts_results + item[ + "chunk_id" + ] + for item + in fts_results } vector_results = [ item - for item in vector_results - if item["chunk_id"] + for item + in vector_results + if item[ + "chunk_id" + ] in fts_chunk_ids ] - hybrid_results = fuse_hybrid_results( - fts_results, - vector_results, + # ------------------------- + # Hybrid fusion + # ------------------------- + + hybrid_results = ( + fuse_hybrid_results( + fts_results, + vector_results, + ) ) + # ------------------------- + # Document diversification + # ------------------------- + + final_results = ( + diversify_results( + hybrid_results, + limit, + max_per_document, + ) + ) + return { "engine": ( "hybrid_fts5_embeddings" @@ -1192,9 +234,7 @@ def search_database( "strategies": ( used_strategies ), - "results": diversify_results( - hybrid_results, - limit, - max_per_document, + "results": ( + final_results ), }