From 0263d4d068c1c2785fbc43a21616c983d1744172 Mon Sep 17 00:00:00 2001
From: Debanjum Singh Solanky <debanjum@gmail.com>
Date: Sun, 29 Aug 2021 03:07:36 -0700
Subject: [PATCH] Enable semantic search for songs in org-music

Org-Music: https://github.com/debanjum/org-music
---
 src/interface/emacs/semantic-search.el |  4 ++-
 src/main.py                            | 38 ++++++++++++++++++++++++++
 src/utils/cli.py                       |  7 ++++-
 3 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/src/interface/emacs/semantic-search.el b/src/interface/emacs/semantic-search.el
index 4ff5be0b..ca2c4aae 100644
--- a/src/interface/emacs/semantic-search.el
+++ b/src/interface/emacs/semantic-search.el
@@ -65,6 +65,7 @@
 (defun semantic-search--buffer-name-to-search-type (buffer-name)
   (let ((file-extension (file-name-extension buffer-name)))
     (cond
+     ((equal buffer-name "Music.org") "music")
      ((equal file-extension "bean") "ledger")
      ((equal file-extension "org") "notes")
      (t "notes"))))
@@ -90,10 +91,11 @@
             (json-response (json-parse-buffer :object-type 'alist)))
         (erase-buffer)
         (insert
-         (cond ((equal search-type "notes") (semantic-search--extract-entries-as-org json-response))
+         (cond ((or (equal search-type "notes") (equal search-type "music")) (semantic-search--extract-entries-as-org json-response))
                ((equal search-type "ledger") (semantic-search--extract-entries-as-ledger json-response))
                (t (format "%s" json-response)))))
       (cond ((equal search-type "notes") (org-mode))
+            ((equal search-type "music") (progn (org-mode) (org-music-mode)))
             (t (fundamental-mode)))
       (read-only-mode t))
     (switch-to-buffer buff)))
diff --git a/src/main.py b/src/main.py
index b8b03e06..258a0237 100644
--- a/src/main.py
+++ b/src/main.py
@@ -38,6 +38,19 @@ def search(q: str, n: Optional[int] = 5, t: Optional[str] = None):
         # collate and return results
         return asymmetric.collate_results(hits, entries, results_count)
 
+    if (t == 'music' or t == None) and music_search_enabled:
+        # query music library
+        hits = asymmetric.query_notes(
+            user_query,
+            song_embeddings,
+            songs,
+            song_encoder,
+            song_cross_encoder,
+            song_top_k)
+
+        # collate and return results
+        return asymmetric.collate_results(hits, songs, results_count)
+
     if (t == 'ledger' or t == None) and ledger_search_enabled:
         # query transactions
         hits = symmetric_ledger.query_transactions(
@@ -84,6 +97,18 @@ def regenerate(t: Optional[str] = None):
             regenerate=True,
             verbose=args.verbose)
 
+    if (t == 'music' or t == None) and music_search_enabled:
+        # Extract Entries, Generate Song Embeddings
+        global song_embeddings
+        global songs
+        songs, song_embeddings, _, _, _ = asymmetric.setup(
+            song_config['input-files'],
+            song_config['input-filter'],
+            pathlib.Path(song_config['compressed-jsonl']),
+            pathlib.Path(song_config['embeddings-file']),
+            regenerate=True,
+            verbose=args.verbose)
+
     if (t == 'ledger' or t == None) and ledger_search_enabled:
         # Extract Entries, Generate Embeddings
         global transaction_embeddings
@@ -125,6 +150,19 @@ if __name__ == '__main__':
             args.regenerate,
             args.verbose)
 
+    # Initialize Org Music Search
+    song_config = get_from_dict(args.config, 'content-type', 'music')
+    music_search_enabled = False
+    if song_config and ('input-files' in song_config or 'input-filter' in song_config):
+        music_search_enabled = True
+        songs, song_embeddings, song_encoder, song_cross_encoder, song_top_k = asymmetric.setup(
+            song_config['input-files'],
+            song_config['input-filter'],
+            pathlib.Path(song_config['compressed-jsonl']),
+            pathlib.Path(song_config['embeddings-file']),
+            args.regenerate,
+            args.verbose)
+
     # Initialize Ledger Search
     ledger_config = get_from_dict(args.config, 'content-type', 'ledger')
     ledger_search_enabled = False
diff --git a/src/utils/cli.py b/src/utils/cli.py
index f61f39c6..106c9d89 100644
--- a/src/utils/cli.py
+++ b/src/utils/cli.py
@@ -57,7 +57,12 @@ default_config = {
         'image':
         {
             'embeddings-file': '.image_embeddings.pt'
-        }
+        },
+        'music':
+        {
+            'compressed-jsonl': '.songs.jsonl.gz',
+            'embeddings-file': '.song_embeddings.pt'
+        },
     },
     'search-type':
     {