diff --git a/sample_config.yml b/sample_config.yml index 2203a694..8d5de409 100644 --- a/sample_config.yml +++ b/sample_config.yml @@ -24,12 +24,19 @@ content-type: embeddings-file: "tests/data/.song_embeddings.pt" search-type: + symmetric: + encoder: "sentence-transformers/paraphrase-MiniLM-L6-v2" + cross-encoder: "cross-encoder/ms-marco-MiniLM-L-6-v2" + model_directory: "tests/data/.symmetric" + asymmetric: encoder: "sentence-transformers/msmarco-MiniLM-L-6-v3" cross-encoder: "cross-encoder/ms-marco-MiniLM-L-6-v2" + model_directory: "tests/data/.asymmetric" image: encoder: "clip-ViT-B-32" + model_directory: "tests/data/.image_encoder" processor: conversation: diff --git a/src/main.py b/src/main.py index dea7cf63..de6aa952 100644 --- a/src/main.py +++ b/src/main.py @@ -130,22 +130,22 @@ def initialize_search(config: FullConfig, regenerate: bool, t: SearchType = None # Initialize Org Notes Search if (t == SearchType.Notes or t == None) and config.content_type.org: # Extract Entries, Generate Notes Embeddings - model.notes_search = asymmetric.setup(config.content_type.org, regenerate=regenerate, verbose=verbose) + model.notes_search = asymmetric.setup(config.content_type.org, search_config=config.search_type.asymmetric, regenerate=regenerate, verbose=verbose) # Initialize Org Music Search if (t == SearchType.Music or t == None) and config.content_type.music: # Extract Entries, Generate Music Embeddings - model.music_search = asymmetric.setup(config.content_type.music, regenerate=regenerate, verbose=verbose) + model.music_search = asymmetric.setup(config.content_type.music, search_config=config.search_type.asymmetric, regenerate=regenerate, verbose=verbose) # Initialize Ledger Search if (t == SearchType.Ledger or t == None) and config.content_type.ledger: # Extract Entries, Generate Ledger Embeddings - model.ledger_search = symmetric_ledger.setup(config.content_type.ledger, regenerate=regenerate, verbose=verbose) + model.ledger_search = symmetric_ledger.setup(config.content_type.ledger, search_config=config.search_type.symmetric, regenerate=regenerate, verbose=verbose) # Initialize Image Search if (t == SearchType.Image or t == None) and config.content_type.image: # Extract Entries, Generate Image Embeddings - model.image_search = image_search.setup(config.content_type.image, regenerate=regenerate, verbose=verbose) + model.image_search = image_search.setup(config.content_type.image, search_config=config.search_type.image, regenerate=regenerate, verbose=verbose) return model diff --git a/src/search_type/asymmetric.py b/src/search_type/asymmetric.py index 416cf7e2..670cc39e 100644 --- a/src/search_type/asymmetric.py +++ b/src/search_type/asymmetric.py @@ -12,18 +12,31 @@ import torch from sentence_transformers import SentenceTransformer, CrossEncoder, util # Internal Packages -from src.utils.helpers import get_absolute_path, resolve_absolute_path +from src.utils.helpers import get_absolute_path, resolve_absolute_path, load_model from src.processor.org_mode.org_to_jsonl import org_to_jsonl from src.utils.config import TextSearchModel -from src.utils.rawconfig import TextSearchConfig +from src.utils.rawconfig import AsymmetricConfig, TextSearchConfig -def initialize_model(): +def initialize_model(search_config: AsymmetricConfig): "Initialize model for assymetric semantic search. That is, where query smaller than results" torch.set_num_threads(4) - bi_encoder = SentenceTransformer('sentence-transformers/msmarco-MiniLM-L-6-v3') # The bi-encoder encodes all entries to use for semantic search - top_k = 30 # Number of entries we want to retrieve with the bi-encoder - cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') # The cross-encoder re-ranks the results to improve quality + + # Number of entries we want to retrieve with the bi-encoder + top_k = 30 + + # The bi-encoder encodes all entries to use for semantic search + bi_encoder = load_model( + model_dir = search_config.model_directory, + model_name = search_config.encoder, + model_type = SentenceTransformer) + + # The cross-encoder re-ranks the results to improve quality + cross_encoder = load_model( + model_dir = search_config.model_directory, + model_name = search_config.cross_encoder, + model_type = CrossEncoder) + return bi_encoder, cross_encoder, top_k @@ -149,9 +162,9 @@ def collate_results(hits, entries, count=5): in hits[0:count]] -def setup(config: TextSearchConfig, regenerate: bool, verbose: bool=False) -> TextSearchModel: +def setup(config: TextSearchConfig, search_config: AsymmetricConfig, regenerate: bool, verbose: bool=False) -> TextSearchModel: # Initialize Model - bi_encoder, cross_encoder, top_k = initialize_model() + bi_encoder, cross_encoder, top_k = initialize_model(search_config) # Map notes in Org-Mode files to (compressed) JSONL formatted file if not resolve_absolute_path(config.compressed_jsonl).exists() or regenerate: diff --git a/src/search_type/image_search.py b/src/search_type/image_search.py index f8af7d8e..9f4d04f2 100644 --- a/src/search_type/image_search.py +++ b/src/search_type/image_search.py @@ -10,16 +10,22 @@ from tqdm import trange import torch # Internal Packages -from src.utils.helpers import resolve_absolute_path +from src.utils.helpers import resolve_absolute_path, load_model import src.utils.exiftool as exiftool from src.utils.config import ImageSearchModel -from src.utils.rawconfig import ImageSearchConfig +from src.utils.rawconfig import ImageSearchConfig, ImageSearchTypeConfig -def initialize_model(): +def initialize_model(search_config: ImageSearchTypeConfig): # Initialize Model torch.set_num_threads(4) - encoder = SentenceTransformer('sentence-transformers/clip-ViT-B-32') #Load the CLIP model + + # Load the CLIP model + encoder = load_model( + model_dir = search_config.model_directory, + model_name = search_config.encoder, + model_type = SentenceTransformer) + return encoder @@ -154,9 +160,9 @@ def collate_results(hits, image_names, image_directory, count=5): in hits[0:count]] -def setup(config: ImageSearchConfig, regenerate: bool, verbose: bool=False) -> ImageSearchModel: +def setup(config: ImageSearchConfig, search_config: ImageSearchTypeConfig, regenerate: bool, verbose: bool=False) -> ImageSearchModel: # Initialize Model - encoder = initialize_model() + encoder = initialize_model(search_config) # Extract Entries image_directory = resolve_absolute_path(config.input_directory, strict=True) diff --git a/src/search_type/symmetric_ledger.py b/src/search_type/symmetric_ledger.py index 4a747215..c7d9cc9a 100644 --- a/src/search_type/symmetric_ledger.py +++ b/src/search_type/symmetric_ledger.py @@ -10,18 +10,31 @@ import torch from sentence_transformers import SentenceTransformer, CrossEncoder, util # Internal Packages -from src.utils.helpers import get_absolute_path, resolve_absolute_path +from src.utils.helpers import get_absolute_path, resolve_absolute_path, load_model from src.processor.ledger.beancount_to_jsonl import beancount_to_jsonl from src.utils.config import TextSearchModel -from src.utils.rawconfig import TextSearchConfig +from src.utils.rawconfig import SymmetricConfig, TextSearchConfig -def initialize_model(): +def initialize_model(search_config: SymmetricConfig): "Initialize model for symmetric semantic search. That is, where query of similar size to results" torch.set_num_threads(4) - bi_encoder = SentenceTransformer('sentence-transformers/paraphrase-MiniLM-L6-v2') # The encoder encodes all entries to use for semantic search - top_k = 30 # Number of entries we want to retrieve with the bi-encoder - cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') # The cross-encoder re-ranks the results to improve quality + + # Number of entries we want to retrieve with the bi-encoder + top_k = 30 + + # The bi-encoder encodes all entries to use for semantic search + bi_encoder = load_model( + model_dir = search_config.model_directory, + model_name = search_config.encoder, + model_type = SentenceTransformer) + + # The cross-encoder re-ranks the results to improve quality + cross_encoder = load_model( + model_dir = search_config.model_directory, + model_name = search_config.cross_encoder, + model_type = CrossEncoder) + return bi_encoder, cross_encoder, top_k @@ -141,9 +154,9 @@ def collate_results(hits, entries, count=5): in hits[0:count]] -def setup(config: TextSearchConfig, regenerate: bool, verbose: bool) -> TextSearchModel: +def setup(config: TextSearchConfig, search_config: SymmetricConfig, regenerate: bool, verbose: bool) -> TextSearchModel: # Initialize Model - bi_encoder, cross_encoder, top_k = initialize_model() + bi_encoder, cross_encoder, top_k = initialize_model(search_config) # Map notes in Org-Mode files to (compressed) JSONL formatted file if not resolve_absolute_path(config.compressed_jsonl).exists() or regenerate: diff --git a/src/utils/cli.py b/src/utils/cli.py index aa58af8d..505628a5 100644 --- a/src/utils/cli.py +++ b/src/utils/cli.py @@ -77,14 +77,22 @@ default_config = { }, 'search-type': { - 'asymmetric': + 'symmetric': + { + 'encoder': "sentence-transformers/paraphrase-MiniLM-L6-v2", + 'cross-encoder': "cross-encoder/ms-marco-MiniLM-L-6-v2", + 'model_directory': None + }, + 'asymmetric': { 'encoder': "sentence-transformers/msmarco-MiniLM-L-6-v3", - 'cross-encoder': "cross-encoder/ms-marco-MiniLM-L-6-v2" + 'cross-encoder': "cross-encoder/ms-marco-MiniLM-L-6-v2", + 'model_directory': None }, 'image': { - 'encoder': "clip-ViT-B-32" + 'encoder': "clip-ViT-B-32", + 'model_directory': None }, }, 'processor': diff --git a/src/utils/helpers.py b/src/utils/helpers.py index e2a3b1fe..ab4cec7c 100644 --- a/src/utils/helpers.py +++ b/src/utils/helpers.py @@ -1,4 +1,6 @@ +# Standard Packages import pathlib +from os.path import join def is_none_or_empty(item): @@ -32,3 +34,20 @@ def merge_dicts(priority_dict, default_dict): if k not in priority_dict: merged_dict[k] = default_dict[k] return merged_dict + + +def load_model(model_name, model_dir, model_type): + "Load model from disk or huggingface" + # Construct model path + model_path = join(model_dir, model_name.replace("/", "_")) if model_dir is not None else None + + # Load model from model_path if it exists there + if model_path is not None and resolve_absolute_path(model_path).exists(): + model = model_type(get_absolute_path(model_path)) + # Else load the model from the model_name + else: + model = model_type(model_name) + if model_path is not None: + model.save(model_path) + + return model \ No newline at end of file diff --git a/src/utils/rawconfig.py b/src/utils/rawconfig.py index b28e9a2d..bbb2de31 100644 --- a/src/utils/rawconfig.py +++ b/src/utils/rawconfig.py @@ -37,15 +37,23 @@ class ContentTypeConfig(ConfigBase): image: Optional[ImageSearchConfig] music: Optional[TextSearchConfig] +class SymmetricConfig(ConfigBase): + encoder: Optional[str] + cross_encoder: Optional[str] + model_directory: Optional[Path] + class AsymmetricConfig(ConfigBase): encoder: Optional[str] cross_encoder: Optional[str] + model_directory: Optional[Path] class ImageSearchTypeConfig(ConfigBase): encoder: Optional[str] + model_directory: Optional[Path] class SearchTypeConfig(ConfigBase): asymmetric: Optional[AsymmetricConfig] + symmetric: Optional[SymmetricConfig] image: Optional[ImageSearchTypeConfig] class ConversationProcessorConfig(ConfigBase): diff --git a/tests/conftest.py b/tests/conftest.py index 48c0cf51..bc71068c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,51 +1,78 @@ # Standard Packages import pytest from pathlib import Path +from src import search_type # Internal Packages from src.search_type import asymmetric, image_search -from src.utils.rawconfig import ContentTypeConfig, ImageSearchConfig, TextSearchConfig +from src.utils.rawconfig import AsymmetricConfig, ContentTypeConfig, ImageSearchConfig, ImageSearchTypeConfig, SearchTypeConfig, SymmetricConfig, TextSearchConfig @pytest.fixture(scope='session') -def model_dir(tmp_path_factory): +def search_config(tmp_path_factory): model_dir = tmp_path_factory.mktemp('data') + search_config = SearchTypeConfig() + + search_config.asymmetric = SymmetricConfig( + encoder = "sentence-transformers/paraphrase-MiniLM-L6-v2", + cross_encoder = "cross-encoder/ms-marco-MiniLM-L-6-v2", + model_directory = model_dir + ) + + search_config.asymmetric = AsymmetricConfig( + encoder = "sentence-transformers/msmarco-MiniLM-L-6-v3", + cross_encoder = "cross-encoder/ms-marco-MiniLM-L-6-v2", + model_directory = model_dir + ) + + search_config.image = ImageSearchTypeConfig( + encoder = "clip-ViT-B-32", + model_directory = model_dir + ) + + return search_config + + +@pytest.fixture(scope='session') +def model_dir(search_config): + model_dir = search_config.asymmetric.model_directory + # Generate Image Embeddings from Test Images - search_config = ContentTypeConfig() - search_config.image = ImageSearchConfig( + content_config = ContentTypeConfig() + content_config.image = ImageSearchConfig( input_directory = 'tests/data', embeddings_file = model_dir.joinpath('.image_embeddings.pt'), batch_size = 10, use_xmp_metadata = False) - image_search.setup(search_config.image, regenerate=False, verbose=True) + image_search.setup(content_config.image, search_config.image, regenerate=False, verbose=True) # Generate Notes Embeddings from Test Notes - search_config.org = TextSearchConfig( + content_config.org = TextSearchConfig( input_files = ['tests/data/main_readme.org', 'tests/data/interface_emacs_readme.org'], input_filter = None, compressed_jsonl = model_dir.joinpath('.notes.jsonl.gz'), embeddings_file = model_dir.joinpath('.note_embeddings.pt')) - asymmetric.setup(search_config.org, regenerate=False, verbose=True) + asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=False, verbose=True) return model_dir @pytest.fixture(scope='session') -def search_config(model_dir): - search_config = ContentTypeConfig() - search_config.org = TextSearchConfig( +def content_config(model_dir): + content_config = ContentTypeConfig() + content_config.org = TextSearchConfig( input_files = ['tests/data/main_readme.org', 'tests/data/interface_emacs_readme.org'], input_filter = None, compressed_jsonl = model_dir.joinpath('.notes.jsonl.gz'), embeddings_file = model_dir.joinpath('.note_embeddings.pt')) - search_config.image = ImageSearchConfig( + content_config.image = ImageSearchConfig( input_directory = 'tests/data', - embeddings_file = 'tests/data/.image_embeddings.pt', + embeddings_file = model_dir.joinpath('.image_embeddings.pt'), batch_size = 10, use_xmp_metadata = False) - return search_config + return content_config diff --git a/tests/test_asymmetric_search.py b/tests/test_asymmetric_search.py index 92535b28..14666002 100644 --- a/tests/test_asymmetric_search.py +++ b/tests/test_asymmetric_search.py @@ -1,14 +1,15 @@ # Internal Packages from src.main import model from src.search_type import asymmetric +from src.utils.rawconfig import ContentTypeConfig, SearchTypeConfig # Test # ---------------------------------------------------------------------------------------------------- -def test_asymmetric_setup(search_config): +def test_asymmetric_setup(content_config: ContentTypeConfig, search_config: SearchTypeConfig): # Act # Regenerate notes embeddings during asymmetric setup - notes_model = asymmetric.setup(search_config.org, regenerate=True) + notes_model = asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=True) # Assert assert len(notes_model.entries) == 10 @@ -16,9 +17,9 @@ def test_asymmetric_setup(search_config): # ---------------------------------------------------------------------------------------------------- -def test_asymmetric_search(search_config): +def test_asymmetric_search(content_config: ContentTypeConfig, search_config: SearchTypeConfig): # Arrange - model.notes_search = asymmetric.setup(search_config.org, regenerate=False) + model.notes_search = asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=False) query = "How to git install application?" # Act diff --git a/tests/test_client.py b/tests/test_client.py index d38ca54d..9b695550 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -9,7 +9,7 @@ import pytest from src.main import app, model, config from src.search_type import asymmetric, image_search from src.utils.helpers import resolve_absolute_path -from src.utils.rawconfig import ContentTypeConfig +from src.utils.rawconfig import ContentTypeConfig, SearchTypeConfig # Arrange @@ -18,55 +18,60 @@ client = TestClient(app) # Test # ---------------------------------------------------------------------------------------------------- -def test_search_with_invalid_search_type(): +def test_search_with_invalid_content_type(): # Arrange user_query = "How to call semantic search from Emacs?" # Act - response = client.get(f"/search?q={user_query}&t=invalid_search_type") + response = client.get(f"/search?q={user_query}&t=invalid_content_type") # Assert assert response.status_code == 422 # ---------------------------------------------------------------------------------------------------- -def test_search_with_valid_search_type(search_config: ContentTypeConfig): +def test_search_with_valid_content_type(content_config: ContentTypeConfig, search_config: SearchTypeConfig): # Arrange - config.content_type = search_config + config.content_type = content_config + config.search_type = search_config + # config.content_type.image = search_config.image - for search_type in ["notes", "ledger", "music", "image"]: + for content_type in ["notes", "ledger", "music", "image"]: # Act - response = client.get(f"/search?q=random&t={search_type}") + response = client.get(f"/search?q=random&t={content_type}") # Assert assert response.status_code == 200 # ---------------------------------------------------------------------------------------------------- -def test_regenerate_with_invalid_search_type(): +def test_regenerate_with_invalid_content_type(): # Act - response = client.get(f"/regenerate?t=invalid_search_type") + response = client.get(f"/regenerate?t=invalid_content_type") # Assert assert response.status_code == 422 # ---------------------------------------------------------------------------------------------------- -def test_regenerate_with_valid_search_type(search_config: ContentTypeConfig): +def test_regenerate_with_valid_content_type(content_config: ContentTypeConfig, search_config: SearchTypeConfig): # Arrange - config.content_type = search_config - for search_type in ["notes", "ledger", "music", "image"]: + config.content_type = content_config + config.search_type = search_config + + for content_type in ["notes", "ledger", "music", "image"]: # Act - response = client.get(f"/regenerate?t={search_type}") + response = client.get(f"/regenerate?t={content_type}") # Assert assert response.status_code == 200 # ---------------------------------------------------------------------------------------------------- @pytest.mark.skip(reason="Flaky test. Search doesn't always return expected image path.") -def test_image_search(search_config: ContentTypeConfig): +def test_image_search(content_config: ContentTypeConfig, search_config: SearchTypeConfig): # Arrange - config.content_type = search_config - model.image_search = image_search.setup(search_config.image, regenerate=False) + config.content_type = content_config + config.search_type = search_config + model.image_search = image_search.setup(content_config.image, search_config.image, regenerate=False) query_expected_image_pairs = [("brown kitten next to fallen plant", "kitten_park.jpg"), ("a horse and dog on a leash", "horse_dog.jpg"), ("A guinea pig eating grass", "guineapig_grass.jpg")] @@ -78,16 +83,16 @@ def test_image_search(search_config: ContentTypeConfig): # Assert assert response.status_code == 200 actual_image = Path(response.json()[0]["Entry"]) - expected_image = resolve_absolute_path(search_config.image.input_directory.joinpath(expected_image_name)) + expected_image = resolve_absolute_path(content_config.image.input_directory.joinpath(expected_image_name)) # Assert assert expected_image == actual_image # ---------------------------------------------------------------------------------------------------- -def test_notes_search(search_config: ContentTypeConfig): +def test_notes_search(content_config: ContentTypeConfig, search_config: SearchTypeConfig): # Arrange - model.notes_search = asymmetric.setup(search_config.org, regenerate=False) + model.notes_search = asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=False) user_query = "How to git install application?" # Act @@ -101,9 +106,9 @@ def test_notes_search(search_config: ContentTypeConfig): # ---------------------------------------------------------------------------------------------------- -def test_notes_search_with_include_filter(search_config: ContentTypeConfig): +def test_notes_search_with_include_filter(content_config: ContentTypeConfig, search_config: SearchTypeConfig): # Arrange - model.notes_search = asymmetric.setup(search_config.org, regenerate=False) + model.notes_search = asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=False) user_query = "How to git install application? +Emacs" # Act @@ -117,9 +122,9 @@ def test_notes_search_with_include_filter(search_config: ContentTypeConfig): # ---------------------------------------------------------------------------------------------------- -def test_notes_search_with_exclude_filter(search_config: ContentTypeConfig): +def test_notes_search_with_exclude_filter(content_config: ContentTypeConfig, search_config: SearchTypeConfig): # Arrange - model.notes_search = asymmetric.setup(search_config.org, regenerate=False) + model.notes_search = asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=False) user_query = "How to git install application? -clone" # Act diff --git a/tests/test_image_search.py b/tests/test_image_search.py index 943d6bec..254bcb7f 100644 --- a/tests/test_image_search.py +++ b/tests/test_image_search.py @@ -5,14 +5,15 @@ import pytest from src.main import model from src.search_type import image_search from src.utils.helpers import resolve_absolute_path +from src.utils.rawconfig import ContentTypeConfig, SearchTypeConfig # Test # ---------------------------------------------------------------------------------------------------- -def test_image_search_setup(search_config): +def test_image_search_setup(content_config: ContentTypeConfig, search_config: SearchTypeConfig): # Act # Regenerate image search embeddings during image setup - image_search_model = image_search.setup(search_config.image, regenerate=True) + image_search_model = image_search.setup(content_config.image, search_config.image, regenerate=True) # Assert assert len(image_search_model.image_names) == 3 @@ -21,9 +22,9 @@ def test_image_search_setup(search_config): # ---------------------------------------------------------------------------------------------------- @pytest.mark.skip(reason="results inconsistent currently") -def test_image_search(search_config): +def test_image_search(content_config: ContentTypeConfig, search_config: SearchTypeConfig): # Arrange - model.image_search = image_search.setup(search_config.image, regenerate=False) + model.image_search = image_search.setup(content_config.image, search_config.image, regenerate=False) query_expected_image_pairs = [("brown kitten next to plant", "kitten_park.jpg"), ("horse and dog in a farm", "horse_dog.jpg"), ("A guinea pig eating grass", "guineapig_grass.jpg")] @@ -38,11 +39,11 @@ def test_image_search(search_config): results = image_search.collate_results( hits, model.image_search.image_names, - search_config.image.input_directory, + content_config.image.input_directory, count=1) actual_image = results[0]["Entry"] - expected_image = resolve_absolute_path(search_config.image.input_directory.joinpath(expected_image_name)) + expected_image = resolve_absolute_path(content_config.image.input_directory.joinpath(expected_image_name)) # Assert assert expected_image == actual_image