financial-assistant

1
from django.http import JsonResponse
2
from langchain.retrievers import EnsembleRetriever
3
from langchain_community.retrievers import BM25Retriever
4
from langchain_community.vectorstores import FAISS
5
from langchain_community.embeddings import HuggingFaceEmbeddings
6
import os
7

8
from sentence_transformers import SentenceTransformer
9
from ..assistant import preprocessing, namedEntities, dataLoader
10

11
model_id = 'intfloat/multilingual-e5-base'
12
model = SentenceTransformer(model_id) 
13

14
def get_preprocessed_texts(query): 
15
    preprocessed_query =  preprocessing.preprocessing(query)
16
    lemmas = preprocessed_query.split(' ')
17
    bank_id, product_id = namedEntities.determine_bank_product(lemmas)
18
    prepocessed_db_texts, indicies = dataLoader.get_db_texts(bank_id, product_id)
19
    
20
    return prepocessed_db_texts
21

22

23
def get_faiss_vectorstore(prepocessed_db_texts):
24
    embedding = HuggingFaceEmbeddings(model_name=model_id)
25
    if os.path.exists('faiss_index'):
26
        faiss_vectorstore = FAISS.load_local('faiss_index', embeddings=embedding, allow_dangerous_deserialization=True)
27
    else:
28
        faiss_vectorstore = FAISS.from_texts(prepocessed_db_texts, embedding)
29
        faiss_vectorstore.save_local('faiss_index')
30
    return faiss_vectorstore
31

32

33
def init_retrievers(prepocessed_db_texts, faiss_vectorstore):
34
    bm25_retriever = BM25Retriever.from_texts(prepocessed_db_texts)
35
    bm25_retriever.k = 5
36
    faiss_retriever = faiss_vectorstore.as_retriever(search_kwargs={"k": 5})
37
    
38
    return bm25_retriever, faiss_retriever
39

40

41
def main(request):
42
    query = request.GET.get('q')    
43
    
44
    prepocessed_db_texts = get_preprocessed_texts(query)
45
    faiss_vectorstore = get_faiss_vectorstore(prepocessed_db_texts) 
46
    
47
    bm25_retriever, faiss_retriever = init_retrievers(prepocessed_db_texts, faiss_vectorstore)    
48
    ensemble_retriever = EnsembleRetriever(
49
        retrievers=[bm25_retriever, faiss_retriever], weights=[0.4, 0.6]
50
    )
51
    
52
    results = ensemble_retriever.get_relevant_documents(query, return_source_documents=True)
53
    
54
    for result in results:
55
      bank, category, title, description, link = result.to_json().get('kwargs').get('page_content').split('|')
56
      response_data = {
57
        "bank": bank,
58
        "category": category,
59
        "title": title,
60
        "answer": description,
61
        "link": link
62
      }
63
      return JsonResponse(response_data)