fastrag

1
import argparse
2
import logging
3
from pathlib import Path
4

5
from haystack.lazy_imports import LazyImport
6

7
from fastrag.utils import init_cls, init_haystack_cls, load_yaml
8

9
logger = logging.getLogger(__name__)
10

11
with LazyImport("Please install Qdrant client by: 'pip install qdrant-haystack'") as qdrant_import:
12
    from qdrant_haystack import QdrantDocumentStore
13

14
if __name__ == "__main__":
15
    parser = argparse.ArgumentParser("Create an index using Qdrant as a backend")
16
    parser.add_argument("--store", type=Path, required=True)
17
    parser.add_argument("--data", type=Path, required=True)
18
    parser.add_argument("--embedder", type=Path, required=True)
19
    parser.add_argument("--batch_size", type=int, required=False)
20

21
    args = parser.parse_args()
22
    qdrant_import.check()
23

24
    store_params = load_yaml(args.store)
25
    data_params = load_yaml(args.data)
26
    emb_params = load_yaml(args.embedder)
27

28
    store_cls = store_params.pop("type")
29
    store = init_haystack_cls(store_cls, store_params)
30
    logger.info("Loaded store backend")
31

32
    data_cls = data_params.pop("type")
33
    data = init_cls(data_cls, data_params)
34
    logger.info("Done loading dataset")
35

36
    logger.info("Indexing documents")
37
    for docs in data:
38
        store.write_documents(docs, batch_size=args.batch_size or 100)
39
    logger.info("Done.")
40

41
    logger.info("Loading Embedder")
42
    emb_cls = emb_params.pop("type")
43
    emb_params["document_store"] = store
44
    emb = init_haystack_cls(emb_cls, emb_params)
45

46
    logger.info("Encoding vectors")
47
    store.update_embeddings(emb, batch_size=emb_params["batch_size"])
48
    logger.info("Done.")
49