sij/khoj
1
0
Fork 0
mirror of https://github.com/khoj-ai/khoj.git synced 2024-11-23 23:48:56 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions 16

Reuse search models across sessions. Merge unused pytest fixtures

Browse source 
- Remove unused model_dir pytest fixture. It was only being used by
  the content_config fixture, not by any tests
- Reuse existing search models downloaded to khoj directory.
  Downloading search models for each pytest sessions seems excessive and
  slows down tests quite a bit
This commit is contained in:
Debanjum Singh Solanky 2022-09-10 14:15:43 +03:00
parent 11917c6ddd
commit 2b58218b56
1 changed files with 14 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
tests/conftest.py
View file

@ -4,6 +4,7 @@ import pytest
# Internal Packages # Internal Packages
from src.search_type import image_search, text_search from src.search_type import image_search, text_search
from src.utils.config import SearchType from src.utils.config import SearchType
from src.utils.helpers import resolve_absolute_path
from src.utils.rawconfig import ContentConfig, TextContentConfig, ImageContentConfig, SearchConfig, TextSearchConfig, ImageSearchConfig 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.processor.org_mode.org_to_jsonl import org_to_jsonl
from src.search_filter.date_filter import DateFilter from src.search_filter.date_filter import DateFilter
@ -12,41 +13,41 @@ from src.search_filter.file_filter import FileFilter
@pytest.fixture(scope='session') @pytest.fixture(scope='session')
def search_config(tmp_path_factory) -> SearchConfig: def search_config() -> SearchConfig:
model_dir = tmp_path_factory.mktemp('data') model_dir = resolve_absolute_path('~/.khoj/search')
model_dir.mkdir(parents=True, exist_ok=True)
search_config = SearchConfig() search_config = SearchConfig()
search_config.symmetric = TextSearchConfig( search_config.symmetric = TextSearchConfig(
encoder = "sentence-transformers/all-MiniLM-L6-v2", encoder = "sentence-transformers/all-MiniLM-L6-v2",
cross_encoder = "cross-encoder/ms-marco-MiniLM-L-6-v2", cross_encoder = "cross-encoder/ms-marco-MiniLM-L-6-v2",
model_directory = model_dir model_directory = model_dir / 'symmetric/'
) )
search_config.asymmetric = TextSearchConfig( search_config.asymmetric = TextSearchConfig(
encoder = "sentence-transformers/multi-qa-MiniLM-L6-cos-v1", encoder = "sentence-transformers/multi-qa-MiniLM-L6-cos-v1",
cross_encoder = "cross-encoder/ms-marco-MiniLM-L-6-v2", cross_encoder = "cross-encoder/ms-marco-MiniLM-L-6-v2",
model_directory = model_dir model_directory = model_dir / 'asymmetric/'
) )
search_config.image = ImageSearchConfig( search_config.image = ImageSearchConfig(
encoder = "sentence-transformers/clip-ViT-B-32", encoder = "sentence-transformers/clip-ViT-B-32",
model_directory = model_dir model_directory = model_dir / 'image/'
) )
return search_config return search_config
@pytest.fixture(scope='session') @pytest.fixture(scope='session')
def model_dir(search_config: SearchConfig): def content_config(tmp_path_factory, search_config: SearchConfig):
model_dir = search_config.asymmetric.model_directory content_dir = tmp_path_factory.mktemp('content')
# Generate Image Embeddings from Test Images # Generate Image Embeddings from Test Images
content_config = ContentConfig() content_config = ContentConfig()
content_config.image = ImageContentConfig( content_config.image = ImageContentConfig(
input_directories = ['tests/data/images'], input_directories = ['tests/data/images'],
embeddings_file = model_dir.joinpath('image_embeddings.pt'), embeddings_file = content_dir.joinpath('image_embeddings.pt'),
batch_size = 10, batch_size = 1,
use_xmp_metadata = False) use_xmp_metadata = False)
image_search.setup(content_config.image, search_config.image, regenerate=False) image_search.setup(content_config.image, search_config.image, regenerate=False)
@ -55,28 +56,10 @@ def model_dir(search_config: SearchConfig):
content_config.org = TextContentConfig( content_config.org = TextContentConfig(
input_files = None, input_files = None,
input_filter = 'tests/data/org/*.org', input_filter = 'tests/data/org/*.org',
compressed_jsonl = model_dir.joinpath('notes.jsonl.gz'), compressed_jsonl = content_dir.joinpath('notes.jsonl.gz'),
embeddings_file = model_dir.joinpath('note_embeddings.pt')) embeddings_file = content_dir.joinpath('note_embeddings.pt'))
filters = [DateFilter(), WordFilter(model_dir, search_type=SearchType.Org), FileFilter()] filters = [DateFilter(), WordFilter(content_dir, search_type=SearchType.Org), FileFilter()]
text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, regenerate=False, filters=filters) 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) -> ContentConfig:
content_config = ContentConfig()
content_config.org = TextContentConfig(
input_files = None,
input_filter = 'tests/data/org/*.org',
compressed_jsonl = model_dir.joinpath('notes.jsonl.gz'),
embeddings_file = model_dir.joinpath('note_embeddings.pt'))
content_config.image = ImageContentConfig(
input_directories = ['tests/data/images'],
embeddings_file = model_dir.joinpath('image_embeddings.pt'),
batch_size = 1,
use_xmp_metadata = False)
return content_config return content_config