File size: 1,607 Bytes
d519be4
bc8698e
 
d519be4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0cc9480
d519be4
 
0cc9480
d519be4
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
# from langchain_huggingface import HuggingFaceEmbeddings
from huggingface_hub import snapshot_download
import logging


class VectorStore:
    def __init__(self, embeddings_model, vs_local_path=None, vs_hf_path=None, number_of_contexts=2):
        
        self.number_of_contexts = number_of_contexts

        logging.info("Loading vectorstore...")        
        
        embeddings = HuggingFaceEmbeddings(model_name=embeddings_model)
        logging.info(f"Loaded embeddings model: {embeddings_model}")


        if vs_hf_path:
            hf_vectorstore = snapshot_download(repo_id=vs_hf_path)
            self.vectore_store = FAISS.load_local(hf_vectorstore, embeddings, allow_dangerous_deserialization=True)
            logging.info(f"Loaded vectorstore from {vs_hf_path}")
        else:
            self.vectore_store = FAISS.load_local(vs_local_path, embeddings, allow_dangerous_deserialization=True)
            logging.info(f"Loaded vectorstore from {vs_local_path}")

    
    def get_context(self, instruction, number_of_contexts=2):

        logging.info(f"Getting context for instruction: {instruction}")
        documentos = self.vectore_store.similarity_search_with_score(instruction, k=self.number_of_contexts)
        return self._beautiful_context(documentos)
    
    def _beautiful_context(self, docs):
        context = ""
        for doc in docs:
            context += doc[0].page_content + "\n\n"

        print("Context: ", context)
        return context[:-1]