86 lines
3.2 KiB
Python
86 lines
3.2 KiB
Python
import json
|
|
import os
|
|
from logging import getLogger
|
|
from typing import Dict, List
|
|
|
|
import faiss
|
|
import pickle
|
|
|
|
|
|
class TrainingContentKnowledgeBase:
|
|
|
|
def __init__(self, embeddings, path: str = 'pathways_2_rw_with_ids.json'):
|
|
self._embedding_model = embeddings
|
|
self._tips = None # self._read_json(path)
|
|
self._category_metadata = None
|
|
self._indices = None
|
|
self._logger = getLogger()
|
|
|
|
@staticmethod
|
|
def _read_json(path: str) -> Dict[str, any]:
|
|
with open(path, 'r', encoding="utf-8") as json_file:
|
|
return json.loads(json_file.read())
|
|
|
|
def print_category_count(self):
|
|
category_tips = {}
|
|
for unit in self._tips['units']:
|
|
for page in unit['pages']:
|
|
for tip in page['tips']:
|
|
category = tip['category'].lower().replace(" ", "_")
|
|
if category not in category_tips:
|
|
category_tips[category] = 0
|
|
else:
|
|
category_tips[category] = category_tips[category] + 1
|
|
print(category_tips)
|
|
|
|
def create_embeddings_and_save_them(self) -> None:
|
|
category_embeddings = {}
|
|
category_metadata = {}
|
|
|
|
for unit in self._tips['units']:
|
|
for page in unit['pages']:
|
|
for tip in page['tips']:
|
|
category = tip['category'].lower().replace(" ", "_")
|
|
if category not in category_embeddings:
|
|
category_embeddings[category] = []
|
|
category_metadata[category] = []
|
|
|
|
category_embeddings[category].append(tip['embedding'])
|
|
category_metadata[category].append({"id": tip['id'], "text": tip['text']})
|
|
|
|
category_indices = {}
|
|
for category, embeddings in category_embeddings.items():
|
|
embeddings_array = self._embedding_model.encode(embeddings)
|
|
index = faiss.IndexFlatL2(embeddings_array.shape[1])
|
|
index.add(embeddings_array)
|
|
category_indices[category] = index
|
|
|
|
faiss.write_index(index, f"./faiss/{category}_tips_index.faiss")
|
|
|
|
with open("./faiss/tips_metadata.pkl", "wb") as f:
|
|
pickle.dump(category_metadata, f)
|
|
|
|
def load_indices_and_metadata(
|
|
self,
|
|
directory: str = './faiss',
|
|
suffix: str = '_tips_index.faiss',
|
|
metadata_path: str = './faiss/tips_metadata.pkl'
|
|
):
|
|
files = os.listdir(directory)
|
|
self._indices = {}
|
|
for file in files:
|
|
if file.endswith(suffix):
|
|
self._indices[file[:-len(suffix)]] = faiss.read_index(f'{directory}/{file}')
|
|
self._logger.info(f'Loaded embeddings for {file[:-len(suffix)]} category.')
|
|
|
|
with open(metadata_path, 'rb') as f:
|
|
self._category_metadata = pickle.load(f)
|
|
self._logger.info("Loaded tips metadata")
|
|
|
|
def query_knowledge_base(self, query: str, category: str, top_k: int = 5) -> List[Dict[str, str]]:
|
|
query_embedding = self._embedding_model.encode([query])
|
|
index = self._indices[category]
|
|
D, I = index.search(query_embedding, top_k)
|
|
results = [self._category_metadata[category][i] for i in I[0]]
|
|
return results
|