| import torch |
| from torch import nn |
| from transformers import AutoTokenizer, AutoModel |
| import os.path |
| import faiss |
|
|
| class EmbeddingRetriever: |
| def __init__(self, embedding_model_name, PATH_IDX, chunks): |
| self.PATH_IDX = PATH_IDX |
| self.embedding_model_name = embedding_model_name |
| self.tokenizer = AutoTokenizer.from_pretrained(embedding_model_name) |
| self.embedding_model = AutoModel.from_pretrained(embedding_model_name) |
|
|
| self.index = self.get_idx(chunks) |
|
|
| def get_idx(self, chunks): |
| if os.path.exists(self.PATH_IDX): |
| index = self.load_faiss_index(self.PATH_IDX) |
| else: |
| encoded_docs = self.tokenizer(["source: {}, content: {}".format(chunk.metadata['source'], chunk.page_content) for chunk in chunks], |
| padding = 'max_length', |
| return_tensors="pt") |
| word_embeddings = self.embedding_model(**encoded_docs).last_hidden_state |
|
|
| index = self.build_faiss_index(word_embeddings) |
| self.save_faiss_index(index, self.PATH_IDX) |
| return index |
|
|
|
|
| def retrieve_data(self, query, TOP_K): |
| query_tokens = self.tokenizer(query, padding = 'max_length', return_tensors="pt") |
| query_embedding = self.embedding_model(**query_tokens).last_hidden_state |
| m = nn.Flatten() |
| np_query_embedding = m(query_embedding).detach().numpy() |
|
|
| distances, indices = self.index.search(np_query_embedding, TOP_K) |
|
|
| return indices[0] |
|
|
| def build_faiss_index(self,embeddings): |
| """Builds a FAISS index for efficient similarity search.""" |
| m = nn.Flatten() |
| embeddings = m(embeddings) |
| dimension = embeddings.shape[1] |
| index = faiss.IndexFlatL2(dimension) |
| np_emb = embeddings.detach().numpy() |
| print("shape index:",np_emb.shape) |
| index.add(np_emb) |
| return index |
|
|
| def save_faiss_index(self, index, index_file_path): |
| """Saves a FAISS index to a file.""" |
| faiss.write_index(index, index_file_path) |
| print(f"FAISS index saved to {index_file_path}") |
|
|
| def load_faiss_index(self, index_file_path): |
| """Loads a FAISS index from a file.""" |
| if os.path.exists(index_file_path): |
| index = faiss.read_index(index_file_path) |
| print(f"FAISS index loaded from {index_file_path}") |
| return index |
| else: |
| return None |
