# AUTOGENERATED! DO NOT EDIT! File to edit: 04_indexer.ipynb (unless otherwise specified).

__all__ = ['join_all', 'build_treemap', 'save_archives']

# Cell
def join_all(db, new_files, new_embeddings):
    start = len(db)
    for i, file in enumerate(new_files):
        path, slug = file
        index = i + start
        db[index] = {
            'slug': slug,
            'fpath': path,
            'embed': new_embeddings[i],
        }
    return(db)

# Cell
from annoy import AnnoyIndex

# Cell
def build_treemap(db):
    treemap = AnnoyIndex(512, 'angular')
    for k, v in db.items():
        treemap.add_item(k, v['embed'])

    # Build the treemap, with 5 trees rn
    treemap.build(5)

    return(treemap)


# Cell
import torch

# Cell
def save_archives(root, treemap, db):
    dbpath = root/'memery.pt'
    if dbpath.exists():
#         dbpath.rename(root/'memery-bak.pt')
        dbpath.unlink()
    torch.save(db, dbpath)

    treepath = root/'memery.ann'
    if treepath.exists():
#         treepath.rename(root/'memery-bak.ann')
        treepath.unlink()
    treemap.save(str(treepath))

    return(str(dbpath), str(treepath))