fastrag

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

5
from tqdm import tqdm
6

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

9
logger = logging.getLogger(__name__)
10

11
if __name__ == "__main__":
12
    parser = argparse.ArgumentParser("Embed data and save to pickled file.")
13
    parser.add_argument("--data", type=Path, required=True)
14
    parser.add_argument("--embedder", type=Path, required=True)
15
    parser.add_argument("--store", type=Path, required=True)
16
    parser.add_argument("--batch_num", type=int, required=False)
17
    parser.add_argument("--batch_size", type=int, required=False)
18

19
    args = parser.parse_args()
20

21
    store_params = load_yaml(args.store)
22
    store_cls = store_params.pop("type")
23
    store = init_haystack_cls(store_cls, store_params)
24
    logger.info("Loaded store backend")
25

26
    data_params = load_yaml(args.data)
27
    emb_params = load_yaml(args.embedder)
28

29
    data_cls = data_params.pop("type")
30
    data = init_cls(data_cls, data_params)
31
    logger.info("Done loading dataset")
32

33
    logger.info("Loading Embedder")
34
    emb_cls = emb_params.pop("type")
35
    emb = init_haystack_cls(emb_cls, emb_params)
36

37
    # Can start at a given batch, if provided
38
    batch_start = args.batch_num or 0
39

40
    logger.info("Creating Embeddings...")
41

42
    for batch_i in tqdm(
43
        list(range(data.chunks)),
44
        desc="Data chunks",
45
    ):
46
        if batch_i >= batch_start:
47
            batch = []
48
            docs = data.process(batch_i)  # requires the HFDatasetLoader
49
            emb_batch = emb.embed_documents(docs)
50
            for d, e in zip(docs, emb_batch):
51
                d.embedding = e
52
                batch.append(d.to_dict())
53
            store.write_documents(batch, batch_size=args.batch_size or 100)
54

55
    logger.info("Done.")
56