From a4e6d4502935c53130d7e3b6162ada7a42e76083 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Wed, 29 Mar 2023 22:47:24 -0400 Subject: [PATCH] Updated dataset pickle. --- src/main.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/src/main.py b/src/main.py index 4024626..4b92ba3 100644 --- a/src/main.py +++ b/src/main.py @@ -27,7 +27,7 @@ if str(src_path) not in sys.path: from dataset import create_dataset #from assistant import semantic_search -from settings import PATH_TO_DATASET, EMBEDDING_MODEL +from settings import PATH_TO_DATASET_PKL, EMBEDDING_MODEL, PATH_TO_DATASET_DICT_PKL import numpy as np @@ -154,5 +154,18 @@ def plot_likelihood(embeddings, num_buckets=200): if __name__ == "__main__": - #load_rawdata_into_pkl() - print_out_dataset_stuff() \ No newline at end of file + # load_rawdata_into_pkl() + # print_out_dataset_stuff() + + with open(PATH_TO_DATASET_PKL, 'rb') as f: + dataset = pickle.load(f) + + dataset_dict = { + "embedding_strings": dataset.embedding_strings, + "embeddings": dataset.embeddings, + "embeddings_metadata_index": dataset.embeddings_metadata_index, + "metadata": dataset.metadata + } + + with open(PATH_TO_DATASET_DICT_PKL, 'wb') as f: + pickle.dump(dataset_dict, f) \ No newline at end of file