mirror of
https://github.com/khoj-ai/khoj.git
synced 2025-02-17 08:04:21 +00:00
Disable image search in tests. Import global state from constants module
- Upstream issues causing load of image search model to fail. Disable tests related to image search for now
This commit is contained in:
parent
b04c84721b
commit
bc423d8f76
4 changed files with 21 additions and 19 deletions
|
@ -6,6 +6,7 @@ import torch
|
|||
from src.search_type import image_search, text_search
|
||||
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.utils import constants
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
|
@ -37,17 +38,16 @@ def search_config(tmp_path_factory):
|
|||
@pytest.fixture(scope='session')
|
||||
def model_dir(search_config):
|
||||
model_dir = search_config.asymmetric.model_directory
|
||||
device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
|
||||
|
||||
# Generate Image Embeddings from Test Images
|
||||
content_config = ContentConfig()
|
||||
content_config.image = ImageContentConfig(
|
||||
input_directories = ['tests/data/images'],
|
||||
embeddings_file = model_dir.joinpath('image_embeddings.pt'),
|
||||
batch_size = 10,
|
||||
use_xmp_metadata = False)
|
||||
# content_config = ContentConfig()
|
||||
# content_config.image = ImageContentConfig(
|
||||
# input_directories = ['tests/data/images'],
|
||||
# embeddings_file = model_dir.joinpath('image_embeddings.pt'),
|
||||
# batch_size = 10,
|
||||
# use_xmp_metadata = False)
|
||||
|
||||
image_search.setup(content_config.image, 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
|
||||
content_config.org = TextContentConfig(
|
||||
|
@ -56,7 +56,7 @@ 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, regenerate=False, device=device, verbose=True)
|
||||
text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, regenerate=False, device=constants.device, verbose=True)
|
||||
|
||||
return model_dir
|
||||
|
||||
|
@ -70,10 +70,10 @@ def content_config(model_dir):
|
|||
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 = 10,
|
||||
use_xmp_metadata = False)
|
||||
# content_config.image = ImageContentConfig(
|
||||
# input_directories = ['tests/data/images'],
|
||||
# embeddings_file = model_dir.joinpath('image_embeddings.pt'),
|
||||
# batch_size = 10,
|
||||
# use_xmp_metadata = False)
|
||||
|
||||
return content_config
|
|
@ -2,7 +2,7 @@
|
|||
from pathlib import Path
|
||||
|
||||
# Internal Packages
|
||||
from src.main import model
|
||||
from src.utils.constants import model
|
||||
from src.search_type import text_search
|
||||
from src.utils.rawconfig import ContentConfig, SearchConfig
|
||||
from src.processor.org_mode.org_to_jsonl import org_to_jsonl
|
||||
|
|
|
@ -7,7 +7,8 @@ from fastapi.testclient import TestClient
|
|||
import pytest
|
||||
|
||||
# Internal Packages
|
||||
from src.main import app, model, config
|
||||
from src.main import app
|
||||
from src.utils.constants 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
|
||||
|
@ -37,7 +38,7 @@ def test_search_with_valid_content_type(content_config: ContentConfig, search_co
|
|||
config.search_type = search_config
|
||||
|
||||
# config.content_type.image = search_config.image
|
||||
for content_type in ["org", "markdown", "ledger", "music", "image"]:
|
||||
for content_type in ["org", "markdown", "ledger", "music"]:
|
||||
# Act
|
||||
response = client.get(f"/search?q=random&t={content_type}")
|
||||
# Assert
|
||||
|
@ -59,7 +60,7 @@ def test_reload_with_valid_content_type(content_config: ContentConfig, search_co
|
|||
config.content_type = content_config
|
||||
config.search_type = search_config
|
||||
|
||||
for content_type in ["org", "markdown", "ledger", "music", "image"]:
|
||||
for content_type in ["org", "markdown", "ledger", "music"]:
|
||||
# Act
|
||||
response = client.get(f"/reload?t={content_type}")
|
||||
# Assert
|
||||
|
|
|
@ -6,7 +6,7 @@ from PIL import Image
|
|||
import pytest
|
||||
|
||||
# Internal Packages
|
||||
from src.main import model, web_directory
|
||||
from src.utils.constants import model, web_directory
|
||||
from src.search_type import image_search
|
||||
from src.utils.helpers import resolve_absolute_path
|
||||
from src.utils.rawconfig import ContentConfig, SearchConfig
|
||||
|
@ -14,6 +14,7 @@ from src.utils.rawconfig import ContentConfig, SearchConfig
|
|||
|
||||
# Test
|
||||
# ----------------------------------------------------------------------------------------------------
|
||||
@pytest.mark.skip(reason="upstream issues in loading image search model. disabled for now")
|
||||
def test_image_search_setup(content_config: ContentConfig, search_config: SearchConfig):
|
||||
# Act
|
||||
# Regenerate image search embeddings during image setup
|
||||
|
|
Loading…
Add table
Reference in a new issue