From b63026d97c86256bb0187ae8ac7c1fce40d9f8da Mon Sep 17 00:00:00 2001 From: Debanjum Singh Solanky Date: Fri, 14 Jan 2022 16:31:55 -0500 Subject: [PATCH 1/5] Save Asymmetric Search Model to Disk - Improve application load time - Remove dependence on internet to startup application and perform semantic search --- sample_config.yml | 1 + src/main.py | 4 ++-- src/search_type/asymmetric.py | 29 +++++++++++++++++++++-------- src/utils/helpers.py | 18 ++++++++++++++++++ src/utils/rawconfig.py | 1 + 5 files changed, 43 insertions(+), 10 deletions(-) diff --git a/sample_config.yml b/sample_config.yml index 2203a694..ee19c8ee 100644 --- a/sample_config.yml +++ b/sample_config.yml @@ -27,6 +27,7 @@ search-type: 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" diff --git a/src/main.py b/src/main.py index dea7cf63..71397486 100644 --- a/src/main.py +++ b/src/main.py @@ -130,12 +130,12 @@ 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: 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/utils/helpers.py b/src/utils/helpers.py index e2a3b1fe..27bb4c3f 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,19 @@ 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("/", "_")) + + # Load model from model_path if it exists there + if 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) + 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..44e4cce1 100644 --- a/src/utils/rawconfig.py +++ b/src/utils/rawconfig.py @@ -40,6 +40,7 @@ class ContentTypeConfig(ConfigBase): class AsymmetricConfig(ConfigBase): encoder: Optional[str] cross_encoder: Optional[str] + model_directory: Optional[Path] class ImageSearchTypeConfig(ConfigBase): encoder: Optional[str] From 934ec233b0e83305bbdee5c9e77b0f9a16c63f2b Mon Sep 17 00:00:00 2001 From: Debanjum Singh Solanky Date: Fri, 14 Jan 2022 16:46:56 -0500 Subject: [PATCH 2/5] Add Search Config for Symmetric Model. Save Model to Disk --- sample_config.yml | 5 +++++ src/main.py | 2 +- src/search_type/symmetric_ledger.py | 29 +++++++++++++++++++++-------- src/utils/rawconfig.py | 6 ++++++ 4 files changed, 33 insertions(+), 9 deletions(-) diff --git a/sample_config.yml b/sample_config.yml index ee19c8ee..3509f038 100644 --- a/sample_config.yml +++ b/sample_config.yml @@ -24,6 +24,11 @@ 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" diff --git a/src/main.py b/src/main.py index 71397486..2edbc209 100644 --- a/src/main.py +++ b/src/main.py @@ -140,7 +140,7 @@ def initialize_search(config: FullConfig, regenerate: bool, t: SearchType = None # 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: 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/rawconfig.py b/src/utils/rawconfig.py index 44e4cce1..00eb2fc7 100644 --- a/src/utils/rawconfig.py +++ b/src/utils/rawconfig.py @@ -37,6 +37,11 @@ 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] @@ -47,6 +52,7 @@ class ImageSearchTypeConfig(ConfigBase): class SearchTypeConfig(ConfigBase): asymmetric: Optional[AsymmetricConfig] + symmetric: Optional[SymmetricConfig] image: Optional[ImageSearchTypeConfig] class ConversationProcessorConfig(ConfigBase): From 510faa19047ae270ceee5cb54d259252c79a81f3 Mon Sep 17 00:00:00 2001 From: Debanjum Singh Solanky Date: Fri, 14 Jan 2022 17:09:18 -0500 Subject: [PATCH 3/5] Save Image Search Model to Disk --- sample_config.yml | 1 + src/main.py | 2 +- src/search_type/image_search.py | 18 ++++++++++++------ src/utils/rawconfig.py | 1 + 4 files changed, 15 insertions(+), 7 deletions(-) diff --git a/sample_config.yml b/sample_config.yml index 3509f038..8d5de409 100644 --- a/sample_config.yml +++ b/sample_config.yml @@ -36,6 +36,7 @@ search-type: 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 2edbc209..de6aa952 100644 --- a/src/main.py +++ b/src/main.py @@ -145,7 +145,7 @@ def initialize_search(config: FullConfig, regenerate: bool, t: SearchType = None # 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/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/utils/rawconfig.py b/src/utils/rawconfig.py index 00eb2fc7..bbb2de31 100644 --- a/src/utils/rawconfig.py +++ b/src/utils/rawconfig.py @@ -49,6 +49,7 @@ class AsymmetricConfig(ConfigBase): class ImageSearchTypeConfig(ConfigBase): encoder: Optional[str] + model_directory: Optional[Path] class SearchTypeConfig(ConfigBase): asymmetric: Optional[AsymmetricConfig] From c64e0c296516f92d93cbd56e779f16dcb74f6cbc Mon Sep 17 00:00:00 2001 From: Debanjum Singh Solanky Date: Fri, 14 Jan 2022 17:32:45 -0500 Subject: [PATCH 4/5] Load model from HuggingFace if model_directory unset in config YAML - Do not save/load the model to/from disk when model_directory unset in config.yml - Add symmetric search default config to cli.py --- src/utils/cli.py | 14 +++++++++++--- src/utils/helpers.py | 7 ++++--- 2 files changed, 15 insertions(+), 6 deletions(-) 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 27bb4c3f..ab4cec7c 100644 --- a/src/utils/helpers.py +++ b/src/utils/helpers.py @@ -39,14 +39,15 @@ def merge_dicts(priority_dict, default_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("/", "_")) + 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 resolve_absolute_path(model_path).exists(): + 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) - model.save(model_path) + if model_path is not None: + model.save(model_path) return model \ No newline at end of file From ed144f798486d2f8525fe65b7ea5d3e9ccd665fd Mon Sep 17 00:00:00 2001 From: Debanjum Singh Solanky Date: Fri, 14 Jan 2022 20:13:14 -0500 Subject: [PATCH 5/5] Setup Search with Search_Config to Fix Tests - Rename pytest fixture search_config to more appropriate content_config - Create search_config pytest fixture - Use search_config where search being setup, used in tests --- tests/conftest.py | 53 +++++++++++++++++++++++++-------- tests/test_asymmetric_search.py | 9 +++--- tests/test_client.py | 51 +++++++++++++++++-------------- tests/test_image_search.py | 13 ++++---- 4 files changed, 80 insertions(+), 46 deletions(-) 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