Katon Minhas
c4e519894b
Refactor/one time io * Delete client work * Streamline imports * Fix list class * Update tests * Clear code funcs test * Fix unit tests * Single read code funcs * Single-read model * Single-load embeddings * Successful E2E test * update unit tests * Update importsg * Reload poetry.lock * remove old exhibit header function * Remove aarete_derived generic function * remove align and format tables * Remove strings to dict * references * Clear code * Move preprocessing_funcs * remove keywords * refactor postprocessing_funcs * Pass unit test - remove qa_qc directory * Black and isort * remove print Approved-by: Alex Galarce
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)
|