diff --git a/src/main.py b/src/main.py index 3c3ce5b5..5adbdd3e 100644 --- a/src/main.py +++ b/src/main.py @@ -1,5 +1,6 @@ # Standard Packages import sys, json, yaml, os +import time from typing import Optional # External Packages @@ -66,50 +67,74 @@ def search(q: str, n: Optional[int] = 5, t: Optional[SearchType] = None): device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu") user_query = q results_count = n + results = {} if (t == SearchType.Org or t == None) and model.orgmode_search: # query org-mode notes - hits, entries = text_search.query(user_query, model.orgmode_search, device=device, filters=[explicit_filter, date_filter]) + query_start = time.time() + hits, entries = text_search.query(user_query, model.orgmode_search, device=device, filters=[explicit_filter, date_filter], verbose=verbose) + query_end = time.time() # collate and return results - return text_search.collate_results(hits, entries, results_count) + collate_start = time.time() + results = text_search.collate_results(hits, entries, results_count) + collate_end = time.time() if (t == SearchType.Music or t == None) and model.music_search: # query music library - hits, entries = text_search.query(user_query, model.music_search, device=device, filters=[explicit_filter, date_filter]) + query_start = time.time() + hits, entries = text_search.query(user_query, model.music_search, device=device, filters=[explicit_filter, date_filter], verbose=verbose) + query_end = time.time() # collate and return results - return text_search.collate_results(hits, entries, results_count) + collate_start = time.time() + results = text_search.collate_results(hits, entries, results_count) + collate_end = time.time() if (t == SearchType.Markdown or t == None) and model.orgmode_search: # query markdown files - hits, entries = text_search.query(user_query, model.markdown_search, device=device, filters=[explicit_filter, date_filter]) + query_start = time.time() + hits, entries = text_search.query(user_query, model.markdown_search, device=device, filters=[explicit_filter, date_filter], verbose=verbose) + query_end = time.time() # collate and return results - return text_search.collate_results(hits, entries, results_count) + collate_start = time.time() + results = text_search.collate_results(hits, entries, results_count) + collate_end = time.time() if (t == SearchType.Ledger or t == None) and model.ledger_search: # query transactions - hits, entries = text_search.query(user_query, model.ledger_search, filters=[explicit_filter, date_filter]) + query_start = time.time() + hits, entries = text_search.query(user_query, model.ledger_search, filters=[explicit_filter, date_filter], verbose=verbose) + query_end = time.time() # collate and return results - return text_search.collate_results(hits, entries, results_count) + collate_start = time.time() + results = text_search.collate_results(hits, entries, results_count) + collate_end = time.time() if (t == SearchType.Image or t == None) and model.image_search: # query images - hits = image_search.query(user_query, results_count, model.image_search) + query_start = time.time() + hits = image_search.query(user_query, results_count, model.image_search, verbose=verbose) output_directory = f'{os.getcwd()}/{web_directory}' + query_end = time.time() # collate and return results - return image_search.collate_results( + collate_start = time.time() + results = image_search.collate_results( hits, image_names=model.image_search.image_names, output_directory=output_directory, static_files_url='/static', count=results_count) + collate_end = time.time() - else: - return {} + if verbose > 1: + print(f"Query took {query_end - query_start:.3f} seconds") + print(f"Collating results took {collate_end - collate_start:.3f} seconds") + + return results @app.get('/reload') diff --git a/src/processor/org_mode/org_to_jsonl.py b/src/processor/org_mode/org_to_jsonl.py index bf147faa..caad7715 100644 --- a/src/processor/org_mode/org_to_jsonl.py +++ b/src/processor/org_mode/org_to_jsonl.py @@ -82,7 +82,7 @@ def convert_org_entries_to_jsonl(entries, verbose=0): continue entry_dict["compiled"] = f'{entry.Heading()}.' - if verbose > 1: + if verbose > 2: print(f"Title: {entry.Heading()}") if entry.Tags(): diff --git a/src/search_type/text_search.py b/src/search_type/text_search.py index 39ae19b8..93c8c344 100644 --- a/src/search_type/text_search.py +++ b/src/search_type/text_search.py @@ -2,6 +2,7 @@ import argparse import pathlib from copy import deepcopy +import time # External Packages import torch @@ -62,38 +63,62 @@ def compute_embeddings(entries, bi_encoder, embeddings_file, regenerate=False, d return corpus_embeddings -def query(raw_query: str, model: TextSearchModel, device='cpu', filters: list = []): +def query(raw_query: str, model: TextSearchModel, device='cpu', filters: list = [], verbose=0): "Search for entries that answer the query" # Copy original embeddings, entries to filter them for query + start = time.time() query = raw_query corpus_embeddings = deepcopy(model.corpus_embeddings) entries = deepcopy(model.entries) + end = time.time() + if verbose > 1: + print(f"Copy Time: {end - start:.3f} seconds") # Filter query, entries and embeddings before semantic search + start = time.time() for filter in filters: query, entries, corpus_embeddings = filter(query, entries, corpus_embeddings) if entries is None or len(entries) == 0: return [], [] + end = time.time() + if verbose > 1: + print(f"Filter Time: {end - start:.3f} seconds") # Encode the query using the bi-encoder + start = time.time() question_embedding = model.bi_encoder.encode([query], convert_to_tensor=True) question_embedding.to(device) question_embedding = util.normalize_embeddings(question_embedding) + end = time.time() + if verbose > 1: + print(f"Query Encode Time: {end - start:.3f} seconds") # Find relevant entries for the query + start = time.time() hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=model.top_k, score_function=util.dot_score)[0] + end = time.time() + if verbose > 1: + print(f"Search Time: {end - start:.3f} seconds") # Score all retrieved entries using the cross-encoder + start = time.time() cross_inp = [[query, entries[hit['corpus_id']]['compiled']] for hit in hits] cross_scores = model.cross_encoder.predict(cross_inp) + end = time.time() + if verbose > 1: + print(f"Cross-Encoder Predict Time: {end - start:.3f} seconds") # Store cross-encoder scores in results dictionary for ranking for idx in range(len(cross_scores)): hits[idx]['cross-score'] = cross_scores[idx] # Order results by cross-encoder score followed by bi-encoder score + start = time.time() hits.sort(key=lambda x: x['score'], reverse=True) # sort by bi-encoder score hits.sort(key=lambda x: x['cross-score'], reverse=True) # sort by cross-encoder score + end = time.time() + if verbose > 1: + print(f"Rank Time: {end - start:.3f} seconds") return hits, entries