Trace query response performance and display timings in verbose mode

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')

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():

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