File size: 787 Bytes
f66560f
 
 
 
fb95c43
f66560f
fb95c43
f66560f
 
 
 
 
6f22fc5
f66560f
 
 
6f22fc5
f66560f
 
 
 
 
 
 
 
 
3d33782
 
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
# retriever and qa_chain function

# HF libraries
from langchain.llms import HuggingFaceHub
from langchain_huggingface import HuggingFaceHubEmbeddings
# vectorestore
from langchain_community.vectorstores import FAISS
# retrieval chain
from langchain.chains import RetrievalQA
# prompt template
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory
from config import EMBEDDING_MODEL


def get_db_retriever(vector_db:str=None):
    embeddings = HuggingFaceHubEmbeddings(repo_id=EMBEDDING_MODEL)

    if not vector_db:
        FAISS_INDEX_PATH='./vectorstore/py-faiss-multi-mpnet-500'
    else:
        FAISS_INDEX_PATH=vector_db
    db = FAISS.load_local(FAISS_INDEX_PATH, embeddings)

    retriever = db.as_retriever()

    return retriever