mirror of
https://github.com/khoj-ai/khoj.git
synced 2025-02-17 08:04:21 +00:00
Update types used in conftest.py
This commit is contained in:
parent
9b16cdbb41
commit
c4cd4b57f1
2 changed files with 13 additions and 17 deletions
|
@ -3,8 +3,8 @@ import pytest
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
# Internal Packages
|
# Internal Packages
|
||||||
from src.utils.config import SearchConfig, TextSearchConfig, ImageSearchConfig
|
|
||||||
from src.search_type import asymmetric, image_search
|
from src.search_type import asymmetric, image_search
|
||||||
|
from src.utils.rawconfig import SearchConfigModel, ImageSearchConfigModel, TextSearchConfigModel
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(scope='session')
|
@pytest.fixture(scope='session')
|
||||||
|
@ -12,23 +12,21 @@ def model_dir(tmp_path_factory):
|
||||||
model_dir = tmp_path_factory.mktemp('data')
|
model_dir = tmp_path_factory.mktemp('data')
|
||||||
|
|
||||||
# Generate Image Embeddings from Test Images
|
# Generate Image Embeddings from Test Images
|
||||||
search_config = SearchConfig()
|
search_config = SearchConfigModel()
|
||||||
search_config.image = ImageSearchConfig(
|
search_config.image = ImageSearchConfigModel(
|
||||||
input_directory = Path('tests/data'),
|
input_directory = Path('tests/data'),
|
||||||
embeddings_file = model_dir.joinpath('.image_embeddings.pt'),
|
embeddings_file = model_dir.joinpath('.image_embeddings.pt'),
|
||||||
batch_size = 10,
|
batch_size = 10,
|
||||||
use_xmp_metadata = False,
|
use_xmp_metadata = False)
|
||||||
verbose = 2)
|
|
||||||
|
|
||||||
image_search.setup(search_config.image, regenerate=False)
|
image_search.setup(search_config.image, regenerate=False)
|
||||||
|
|
||||||
# Generate Notes Embeddings from Test Notes
|
# Generate Notes Embeddings from Test Notes
|
||||||
search_config.notes = TextSearchConfig(
|
search_config.org = TextSearchConfigModel(
|
||||||
input_files = [Path('tests/data/main_readme.org'), Path('tests/data/interface_emacs_readme.org')],
|
input_files = [Path('tests/data/main_readme.org'), Path('tests/data/interface_emacs_readme.org')],
|
||||||
input_filter = None,
|
input_filter = None,
|
||||||
compressed_jsonl = model_dir.joinpath('.notes.jsonl.gz'),
|
compressed_jsonl = model_dir.joinpath('.notes.jsonl.gz'),
|
||||||
embeddings_file = model_dir.joinpath('.note_embeddings.pt'),
|
embeddings_file = model_dir.joinpath('.note_embeddings.pt'))
|
||||||
verbose = 0)
|
|
||||||
|
|
||||||
asymmetric.setup(search_config.notes, regenerate=False)
|
asymmetric.setup(search_config.notes, regenerate=False)
|
||||||
|
|
||||||
|
@ -37,19 +35,17 @@ def model_dir(tmp_path_factory):
|
||||||
|
|
||||||
@pytest.fixture(scope='session')
|
@pytest.fixture(scope='session')
|
||||||
def search_config(model_dir):
|
def search_config(model_dir):
|
||||||
search_config = SearchConfig()
|
search_config = SearchConfigModel()
|
||||||
search_config.notes = TextSearchConfig(
|
search_config.org = TextSearchConfigModel(
|
||||||
input_files = [Path('tests/data/main_readme.org'), Path('tests/data/interface_emacs_readme.org')],
|
input_files = [Path('tests/data/main_readme.org'), Path('tests/data/interface_emacs_readme.org')],
|
||||||
input_filter = None,
|
input_filter = None,
|
||||||
compressed_jsonl = model_dir.joinpath('.notes.jsonl.gz'),
|
compressed_jsonl = model_dir.joinpath('.notes.jsonl.gz'),
|
||||||
embeddings_file = model_dir.joinpath('.note_embeddings.pt'),
|
embeddings_file = model_dir.joinpath('.note_embeddings.pt'))
|
||||||
verbose = 2)
|
|
||||||
|
|
||||||
search_config.image = ImageSearchConfig(
|
search_config.image = ImageSearchConfigModel(
|
||||||
input_directory = Path('tests/data'),
|
input_directory = Path('tests/data'),
|
||||||
embeddings_file = Path('tests/data/.image_embeddings.pt'),
|
embeddings_file = Path('tests/data/.image_embeddings.pt'),
|
||||||
batch_size = 10,
|
batch_size = 10,
|
||||||
use_xmp_metadata = False,
|
use_xmp_metadata = False)
|
||||||
verbose = 2)
|
|
||||||
|
|
||||||
return search_config
|
return search_config
|
||||||
|
|
|
@ -8,7 +8,7 @@ from src.search_type import asymmetric
|
||||||
def test_asymmetric_setup(search_config):
|
def test_asymmetric_setup(search_config):
|
||||||
# Act
|
# Act
|
||||||
# Regenerate notes embeddings during asymmetric setup
|
# Regenerate notes embeddings during asymmetric setup
|
||||||
notes_model = asymmetric.setup(search_config.notes, regenerate=True)
|
notes_model = asymmetric.setup(search_config.org, regenerate=True)
|
||||||
|
|
||||||
# Assert
|
# Assert
|
||||||
assert len(notes_model.entries) == 10
|
assert len(notes_model.entries) == 10
|
||||||
|
@ -18,7 +18,7 @@ def test_asymmetric_setup(search_config):
|
||||||
# ----------------------------------------------------------------------------------------------------
|
# ----------------------------------------------------------------------------------------------------
|
||||||
def test_asymmetric_search(search_config):
|
def test_asymmetric_search(search_config):
|
||||||
# Arrange
|
# Arrange
|
||||||
model.notes_search = asymmetric.setup(search_config.notes, regenerate=False)
|
model.notes_search = asymmetric.setup(search_config.org, regenerate=False)
|
||||||
query = "How to git install application?"
|
query = "How to git install application?"
|
||||||
|
|
||||||
# Act
|
# Act
|
||||||
|
|
Loading…
Add table
Reference in a new issue