utilities/retrievers.py

import pickle
from logging import getLogger
from pathlib import Path

import torch
from haystack.nodes import EmbeddingRetriever
from qdrant_haystack import QdrantDocumentStore

logger = getLogger(__name__)

proj_dir = Path(__file__).parents[1]

st_document_store_path = proj_dir / 'haystack_pickles' / 'simple-wiki_all-mpnet-base-v2_document-store.pkl'

logger.warning('Loading Document Store...')
with open(st_document_store_path, 'rb') as handle:
    st_document_store = pickle.load(handle)
logger.warning('Loaded Document Store...')

qd_document_store = QdrantDocumentStore(path=str(proj_dir/'Qdrant'))

qd_document_store.main_device = torch.device('cpu')
qd_retriever = EmbeddingRetriever(document_store=qd_document_store,
                                  embedding_model="sentence-transformers/all-mpnet-base-v2",
                                  model_format="sentence_transformers",
                                  use_gpu=True)