From 092b9e329d88c3645ef737d04c2233d181e51007 Mon Sep 17 00:00:00 2001 From: Debanjum Singh Solanky Date: Mon, 5 Sep 2022 01:05:13 +0300 Subject: [PATCH] Setup Filters when configuring Text Search for each Search Type - Allows enabling different filters for different Text Search Types - Use FileFilter in Text Search on Org Files --- src/configure.py | 13 +++++++++---- src/search_type/text_search.py | 9 +++------ tests/conftest.py | 14 +++++++++----- tests/test_client.py | 12 +++++++----- tests/test_text_search.py | 10 +++++----- 5 files changed, 33 insertions(+), 25 deletions(-) diff --git a/src/configure.py b/src/configure.py index 938062eb..3ee594b4 100644 --- a/src/configure.py +++ b/src/configure.py @@ -14,6 +14,9 @@ from src.utils.config import SearchType, SearchModels, ProcessorConfigModel, Con from src.utils import state from src.utils.helpers import resolve_absolute_path from src.utils.rawconfig import FullConfig, ProcessorConfig +from src.search_filter.date_filter import DateFilter +from src.search_filter.word_filter import WordFilter +from src.search_filter.file_filter import FileFilter logger = logging.getLogger(__name__) @@ -39,23 +42,25 @@ def configure_server(args, required=False): def configure_search(model: SearchModels, config: FullConfig, regenerate: bool, t: SearchType = None): # Initialize Org Notes Search if (t == SearchType.Org or t == None) and config.content_type.org: + filter_directory = resolve_absolute_path(config.content_type.org.compressed_jsonl.parent) + filters = [DateFilter(), WordFilter(filter_directory, search_type=SearchType.Org), FileFilter()] # Extract Entries, Generate Notes Embeddings - model.orgmode_search = text_search.setup(org_to_jsonl, config.content_type.org, search_config=config.search_type.asymmetric, search_type=SearchType.Org, regenerate=regenerate) + model.orgmode_search = text_search.setup(org_to_jsonl, config.content_type.org, search_config=config.search_type.asymmetric, regenerate=regenerate, filters=filters) # Initialize Org Music Search if (t == SearchType.Music or t == None) and config.content_type.music: # Extract Entries, Generate Music Embeddings - model.music_search = text_search.setup(org_to_jsonl, config.content_type.music, search_config=config.search_type.asymmetric, search_type=SearchType.Music, regenerate=regenerate) + model.music_search = text_search.setup(org_to_jsonl, config.content_type.music, search_config=config.search_type.asymmetric, regenerate=regenerate) # Initialize Markdown Search if (t == SearchType.Markdown or t == None) and config.content_type.markdown: # Extract Entries, Generate Markdown Embeddings - model.markdown_search = text_search.setup(markdown_to_jsonl, config.content_type.markdown, search_config=config.search_type.asymmetric, search_type=SearchType.Markdown, regenerate=regenerate) + model.markdown_search = text_search.setup(markdown_to_jsonl, config.content_type.markdown, search_config=config.search_type.asymmetric, regenerate=regenerate) # Initialize Ledger Search if (t == SearchType.Ledger or t == None) and config.content_type.ledger: # Extract Entries, Generate Ledger Embeddings - model.ledger_search = text_search.setup(beancount_to_jsonl, config.content_type.ledger, search_config=config.search_type.symmetric, search_type=SearchType.Ledger, regenerate=regenerate) + model.ledger_search = text_search.setup(beancount_to_jsonl, config.content_type.ledger, search_config=config.search_type.symmetric, regenerate=regenerate) # Initialize Image Search if (t == SearchType.Image or t == None) and config.content_type.image: diff --git a/src/search_type/text_search.py b/src/search_type/text_search.py index a674d712..1dc30ef2 100644 --- a/src/search_type/text_search.py +++ b/src/search_type/text_search.py @@ -7,13 +7,12 @@ import time # External Packages import torch from sentence_transformers import SentenceTransformer, CrossEncoder, util -from src.search_filter.date_filter import DateFilter -from src.search_filter.word_filter import WordFilter +from src.search_filter.base_filter import BaseFilter # Internal Packages from src.utils import state from src.utils.helpers import get_absolute_path, resolve_absolute_path, load_model -from src.utils.config import SearchType, TextSearchModel +from src.utils.config import TextSearchModel from src.utils.rawconfig import TextSearchConfig, TextContentConfig from src.utils.jsonl import load_jsonl @@ -153,7 +152,7 @@ def collate_results(hits, entries, count=5): in hits[0:count]] -def setup(text_to_jsonl, config: TextContentConfig, search_config: TextSearchConfig, search_type: SearchType, regenerate: bool) -> TextSearchModel: +def setup(text_to_jsonl, config: TextContentConfig, search_config: TextSearchConfig, regenerate: bool, filters: list[BaseFilter] = []) -> TextSearchModel: # Initialize Model bi_encoder, cross_encoder, top_k = initialize_model(search_config) @@ -170,8 +169,6 @@ def setup(text_to_jsonl, config: TextContentConfig, search_config: TextSearchCon config.embeddings_file = resolve_absolute_path(config.embeddings_file) corpus_embeddings = compute_embeddings(entries, bi_encoder, config.embeddings_file, regenerate=regenerate) - filter_directory = resolve_absolute_path(config.compressed_jsonl.parent) - filters = [DateFilter(), WordFilter(filter_directory, search_type=search_type)] for filter in filters: filter.load(entries, regenerate=regenerate) diff --git a/tests/conftest.py b/tests/conftest.py index 930ec734..7545527f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,4 @@ -# Standard Packages +# External Packages import pytest # Internal Packages @@ -6,10 +6,13 @@ from src.search_type import image_search, text_search from src.utils.config import SearchType from src.utils.rawconfig import ContentConfig, TextContentConfig, ImageContentConfig, SearchConfig, TextSearchConfig, ImageSearchConfig from src.processor.org_mode.org_to_jsonl import org_to_jsonl +from src.search_filter.date_filter import DateFilter +from src.search_filter.word_filter import WordFilter +from src.search_filter.file_filter import FileFilter @pytest.fixture(scope='session') -def search_config(tmp_path_factory): +def search_config(tmp_path_factory) -> SearchConfig: model_dir = tmp_path_factory.mktemp('data') search_config = SearchConfig() @@ -35,7 +38,7 @@ def search_config(tmp_path_factory): @pytest.fixture(scope='session') -def model_dir(search_config): +def model_dir(search_config: SearchConfig): model_dir = search_config.asymmetric.model_directory # Generate Image Embeddings from Test Images @@ -55,13 +58,14 @@ def model_dir(search_config): compressed_jsonl = model_dir.joinpath('notes.jsonl.gz'), embeddings_file = model_dir.joinpath('note_embeddings.pt')) - text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, SearchType.Org, regenerate=False) + filters = [DateFilter(), WordFilter(model_dir, search_type=SearchType.Org), FileFilter()] + text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, regenerate=False, filters=filters) return model_dir @pytest.fixture(scope='session') -def content_config(model_dir): +def content_config(model_dir) -> ContentConfig: content_config = ContentConfig() content_config.org = TextContentConfig( input_files = None, diff --git a/tests/test_client.py b/tests/test_client.py index e7ddac33..578c789c 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -4,7 +4,6 @@ from PIL import Image # External Packages from fastapi.testclient import TestClient -import pytest # Internal Packages from src.main import app @@ -12,7 +11,8 @@ from src.utils.config import SearchType from src.utils.state import model, config from src.search_type import text_search, image_search from src.utils.rawconfig import ContentConfig, SearchConfig -from src.processor.org_mode import org_to_jsonl +from src.processor.org_mode.org_to_jsonl import org_to_jsonl +from src.search_filter.word_filter import WordFilter # Arrange @@ -116,7 +116,7 @@ def test_image_search(content_config: ContentConfig, search_config: SearchConfig # ---------------------------------------------------------------------------------------------------- def test_notes_search(content_config: ContentConfig, search_config: SearchConfig): # Arrange - model.orgmode_search = text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, SearchType.Org, regenerate=False) + model.orgmode_search = text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, regenerate=False) user_query = "How to git install application?" # Act @@ -132,7 +132,8 @@ def test_notes_search(content_config: ContentConfig, search_config: SearchConfig # ---------------------------------------------------------------------------------------------------- def test_notes_search_with_include_filter(content_config: ContentConfig, search_config: SearchConfig): # Arrange - model.orgmode_search = text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, SearchType.Org, regenerate=False) + filters = [WordFilter(content_config.org.compressed_jsonl.parent, search_type=SearchType.Org)] + model.orgmode_search = text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, regenerate=False, filters=filters) user_query = 'How to git install application? +"Emacs"' # Act @@ -148,7 +149,8 @@ def test_notes_search_with_include_filter(content_config: ContentConfig, search_ # ---------------------------------------------------------------------------------------------------- def test_notes_search_with_exclude_filter(content_config: ContentConfig, search_config: SearchConfig): # Arrange - model.orgmode_search = text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, SearchType.Org, regenerate=False) + filters = [WordFilter(content_config.org.compressed_jsonl.parent, search_type=SearchType.Org)] + model.orgmode_search = text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, regenerate=False, filters=filters) user_query = 'How to git install application? -"clone"' # Act diff --git a/tests/test_text_search.py b/tests/test_text_search.py index 84f16df5..d56d304d 100644 --- a/tests/test_text_search.py +++ b/tests/test_text_search.py @@ -14,7 +14,7 @@ from src.processor.org_mode.org_to_jsonl import org_to_jsonl def test_asymmetric_setup(content_config: ContentConfig, search_config: SearchConfig): # Act # Regenerate notes embeddings during asymmetric setup - notes_model = text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, SearchType.Org, regenerate=True) + notes_model = text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, regenerate=True) # Assert assert len(notes_model.entries) == 10 @@ -24,7 +24,7 @@ def test_asymmetric_setup(content_config: ContentConfig, search_config: SearchCo # ---------------------------------------------------------------------------------------------------- def test_asymmetric_search(content_config: ContentConfig, search_config: SearchConfig): # Arrange - model.notes_search = text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, SearchType.Org, regenerate=False) + model.notes_search = text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, regenerate=False) query = "How to git install application?" # Act @@ -47,7 +47,7 @@ def test_asymmetric_search(content_config: ContentConfig, search_config: SearchC # ---------------------------------------------------------------------------------------------------- def test_asymmetric_reload(content_config: ContentConfig, search_config: SearchConfig): # Arrange - initial_notes_model= text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, SearchType.Org, regenerate=False) + initial_notes_model= text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, regenerate=False) assert len(initial_notes_model.entries) == 10 assert len(initial_notes_model.corpus_embeddings) == 10 @@ -60,11 +60,11 @@ def test_asymmetric_reload(content_config: ContentConfig, search_config: SearchC f.write("\n* A Chihuahua doing Tango\n- Saw a super cute video of a chihuahua doing the Tango on Youtube\n") # regenerate notes jsonl, model embeddings and model to include entry from new file - regenerated_notes_model = text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, SearchType.Org, regenerate=True) + regenerated_notes_model = text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, regenerate=True) # Act # reload embeddings, entries, notes model from previously generated notes jsonl and model embeddings files - initial_notes_model = text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, SearchType.Org, regenerate=False) + initial_notes_model = text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, regenerate=False) # Assert assert len(regenerated_notes_model.entries) == 11