financial-assistant

1
import numpy as np, csv
2
from sentence_transformers import SentenceTransformer, util
3
from rank_bm25 import BM25Okapi
4

5
model_id = 'intfloat/multilingual-e5-base'
6
model = SentenceTransformer(model_id) 
7

8
def DenseRetriever(query, vectors):
9
    top_k = 20
10
    vec_query = model.encode(query, normalize_embeddings=True)
11
    sim = util.pytorch_cos_sim(vec_query, vectors)[0]    
12
    top_indices = np.argsort(-sim)[:top_k].tolist()
13
        
14
    return top_indices
15

16
def SparseRetriever(query, filtered_indices):
17
    tokenized_corpus = []
18
    texts = []
19
      
20
    with open('db_texts.csv', 'r', encoding='utf-8') as csvfile:
21
      reader = csv.DictReader(csvfile)
22
      for i, row in enumerate(reader):
23
        texts.append(f"{i} {row['preprocessed']}")
24

25
    for index in filtered_indices:
26
      tokenized_corpus.append(texts[index].split(' '))
27

28
    bm25 = BM25Okapi(tokenized_corpus)
29
    tokenized_query = query.split(" ")
30
    top_results = bm25.get_top_n(tokenized_query, tokenized_corpus, n=20)
31
    top_indices = [int(text[0]) for text in top_results]
32
    
33
    return top_indices