2021-10-03 02:46:29 +00:00
|
|
|
# Standard Packages
|
|
|
|
import pytest
|
2022-06-29 20:59:57 +00:00
|
|
|
import torch
|
2021-10-03 02:46:29 +00:00
|
|
|
|
|
|
|
# Internal Packages
|
2022-07-21 14:05:43 +00:00
|
|
|
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
|
2021-10-03 02:46:29 +00:00
|
|
|
|
|
|
|
|
|
|
|
@pytest.fixture(scope='session')
|
2022-01-15 01:13:14 +00:00
|
|
|
def search_config(tmp_path_factory):
|
2021-10-03 02:46:29 +00:00
|
|
|
model_dir = tmp_path_factory.mktemp('data')
|
|
|
|
|
2022-01-15 01:54:38 +00:00
|
|
|
search_config = SearchConfig()
|
2022-01-15 01:13:14 +00:00
|
|
|
|
2022-07-21 14:05:43 +00:00
|
|
|
search_config.symmetric = TextSearchConfig(
|
2022-07-18 16:16:40 +00:00
|
|
|
encoder = "sentence-transformers/all-MiniLM-L6-v2",
|
2022-01-15 01:13:14 +00:00
|
|
|
cross_encoder = "cross-encoder/ms-marco-MiniLM-L-6-v2",
|
|
|
|
model_directory = model_dir
|
|
|
|
)
|
|
|
|
|
2022-07-21 14:05:43 +00:00
|
|
|
search_config.asymmetric = TextSearchConfig(
|
2022-07-18 16:00:19 +00:00
|
|
|
encoder = "sentence-transformers/multi-qa-MiniLM-L6-cos-v1",
|
2022-01-15 01:13:14 +00:00
|
|
|
cross_encoder = "cross-encoder/ms-marco-MiniLM-L-6-v2",
|
|
|
|
model_directory = model_dir
|
|
|
|
)
|
|
|
|
|
2022-01-15 01:54:38 +00:00
|
|
|
search_config.image = ImageSearchConfig(
|
2022-08-04 19:54:39 +00:00
|
|
|
encoder = "sentence-transformers/clip-ViT-B-32",
|
2022-01-15 01:13:14 +00:00
|
|
|
model_directory = model_dir
|
|
|
|
)
|
|
|
|
|
|
|
|
return search_config
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.fixture(scope='session')
|
|
|
|
def model_dir(search_config):
|
|
|
|
model_dir = search_config.asymmetric.model_directory
|
2022-06-29 20:59:57 +00:00
|
|
|
device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
|
2022-01-15 01:13:14 +00:00
|
|
|
|
2021-10-03 02:46:29 +00:00
|
|
|
# Generate Image Embeddings from Test Images
|
2022-01-15 01:54:38 +00:00
|
|
|
content_config = ContentConfig()
|
|
|
|
content_config.image = ImageContentConfig(
|
2022-07-19 22:54:03 +00:00
|
|
|
input_directories = ['tests/data/images'],
|
2022-01-29 06:57:08 +00:00
|
|
|
embeddings_file = model_dir.joinpath('image_embeddings.pt'),
|
2021-10-03 02:46:29 +00:00
|
|
|
batch_size = 10,
|
2021-12-04 17:02:19 +00:00
|
|
|
use_xmp_metadata = False)
|
2021-10-03 02:46:29 +00:00
|
|
|
|
2022-01-15 01:13:14 +00:00
|
|
|
image_search.setup(content_config.image, search_config.image, regenerate=False, verbose=True)
|
2021-10-03 02:46:29 +00:00
|
|
|
|
|
|
|
# Generate Notes Embeddings from Test Notes
|
2022-01-15 01:54:38 +00:00
|
|
|
content_config.org = TextContentConfig(
|
2022-01-29 06:57:08 +00:00
|
|
|
input_files = None,
|
2022-08-04 19:29:57 +00:00
|
|
|
input_filter = 'tests/data/org/*.org',
|
2022-01-29 06:57:08 +00:00
|
|
|
compressed_jsonl = model_dir.joinpath('notes.jsonl.gz'),
|
|
|
|
embeddings_file = model_dir.joinpath('note_embeddings.pt'))
|
2021-10-03 02:46:29 +00:00
|
|
|
|
2022-07-21 14:05:43 +00:00
|
|
|
text_search.setup(org_to_jsonl, content_config.org, search_config.asymmetric, regenerate=False, device=device, verbose=True)
|
2021-10-03 02:46:29 +00:00
|
|
|
|
|
|
|
return model_dir
|
2021-10-03 03:28:33 +00:00
|
|
|
|
|
|
|
|
|
|
|
@pytest.fixture(scope='session')
|
2022-01-15 01:13:14 +00:00
|
|
|
def content_config(model_dir):
|
2022-01-15 01:54:38 +00:00
|
|
|
content_config = ContentConfig()
|
|
|
|
content_config.org = TextContentConfig(
|
2022-01-29 06:57:08 +00:00
|
|
|
input_files = None,
|
2022-08-04 19:29:57 +00:00
|
|
|
input_filter = 'tests/data/org/*.org',
|
2022-01-29 06:57:08 +00:00
|
|
|
compressed_jsonl = model_dir.joinpath('notes.jsonl.gz'),
|
|
|
|
embeddings_file = model_dir.joinpath('note_embeddings.pt'))
|
2021-10-03 03:28:33 +00:00
|
|
|
|
2022-01-15 01:54:38 +00:00
|
|
|
content_config.image = ImageContentConfig(
|
2022-07-19 22:54:03 +00:00
|
|
|
input_directories = ['tests/data/images'],
|
2022-01-29 06:57:08 +00:00
|
|
|
embeddings_file = model_dir.joinpath('image_embeddings.pt'),
|
2021-10-03 03:28:33 +00:00
|
|
|
batch_size = 10,
|
2021-12-04 17:02:19 +00:00
|
|
|
use_xmp_metadata = False)
|
2021-10-03 03:28:33 +00:00
|
|
|
|
2022-01-29 06:57:08 +00:00
|
|
|
return content_config
|