23 lines
631 B
Python
23 lines
631 B
Python
|
import os
|
||
|
|
import pickle
|
||
|
|
|
||
|
|
import faiss
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
|
||
|
|
class CodeEmbedding:
|
||
|
|
"""Class to manage Embeddings for a Doczy.ai execution"""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
level,
|
||
|
|
index_name="faiss_index.bin",
|
||
|
|
embedding_name="embeddings.npy",
|
||
|
|
choices_name="choices.pkl",
|
||
|
|
):
|
||
|
|
self.level = level
|
||
|
|
self.index = faiss.read_index(os.path.join("embeddings", level, index_name))
|
||
|
|
self.embeddings = np.load(os.path.join("embeddings", level, embedding_name))
|
||
|
|
with open(os.path.join("embeddings", level, choices_name), "rb") as f:
|
||
|
|
self.choices = pickle.load(f)
|