klbr/khoj
1
0
Fork 0
mirror of https://github.com/khoaliber/khoj.git synced 2026-03-02 21:19:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

Setup Search with Search_Config to Fix Tests

Browse Source 
- 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
This commit is contained in:
Debanjum Singh Solanky
2022-01-14 20:13:14 -05:00
parent c64e0c2965
commit ed144f7984
4 changed files with 80 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
tests/conftest.py
View File

@@ -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