Upload 17 files

app.py

@@ -1,5 +1,5 @@
 import re
-
+import numpy as np
 import openai
 import streamlit_scrollable_textbox as stx
 
@@ -8,23 +8,27 @@ import streamlit as st
 
 st.set_page_config(layout="wide")  # isort: split
 
+from utils import nltkmodules
 from utils.entity_extraction import (
     clean_entities,
+    extract_keywords,
     extract_quarter_year,
     extract_ticker_spacy,
     format_entities_flan_alpaca,
     generate_alpaca_ner_prompt,
-    extract_keywords
 )
 from utils.models import (
     generate_entities_flan_alpaca_checkpoint,
     generate_entities_flan_alpaca_inference_api,
     generate_text_flan_t5,
-    get_data,
     get_alpaca_model,
+    get_data,
     get_flan_alpaca_xl_model,
     get_flan_t5_model,
     get_instructor_embedding_model,
+    get_instructor_embedding_model_api,
+    get_bm25_model,
+    preprocess_text,
     get_mpnet_embedding_model,
     get_sgpt_embedding_model,
     get_spacy_model,
@@ -55,6 +59,7 @@ from utils.retriever import (
     sentence_id_combine,
     text_lookup,
     year_quarter_range,
+    get_bm25_search_hits,
 )
 from utils.transcript_retrieval import retrieve_transcript
 from utils.vector_index import (
@@ -62,7 +67,6 @@ from utils.vector_index import (
     create_sparse_embeddings,
     hybrid_score_norm,
 )
-from utils import nltkmodules
 
 st.title("Question Answering on Earnings Call Transcripts")
 
@@ -75,6 +79,8 @@ col1, col2 = st.columns([3, 3], gap="medium")
 
 
 with st.sidebar:
+    use_bm25 = st.checkbox("Use BM25 for filtering results")
+
     ner_choice = st.selectbox("Select NER Model", ["Spacy", "Alpaca"])
     document_type = st.selectbox(
         "Select Query Type", ["Single-Document", "Multi-Document"]
@@ -85,6 +91,18 @@ with st.sidebar:
             ["Single-Company", "Compare Companies"],
         )
 
+
+corpus, bm25 = get_bm25_model(data)
+
+tokenized_query = preprocess_text(query_text).split()
+sparse_scores = np.argsort(bm25.get_scores(tokenized_query), axis=0)[::-1]
+indices_hits = get_bm25_search_hits(corpus, sparse_scores, 50)
+
+if use_bm25 == True:
+    indices = indices_hits
+else:
+    indices = None
+
 if ner_choice == "Spacy":
     ner_model = get_spacy_model()
 
@@ -305,7 +323,7 @@ elif encoder_model == "Instructor":
     )
     pinecone_index_name = "week13-instructor-xl"
     pinecone_index = pinecone.Index(pinecone_index_name)
-    retriever_model = get_instructor_embedding_model()
+    retriever_model = get_instructor_embedding_model_api()
     instruction = (
         "Represent the financial question for retrieving supporting documents:"
     )
@@ -318,7 +336,7 @@ elif encoder_model == "Hybrid Instructor - SPLADE":
     )
     pinecone_index_name = "week13-splade-instructor-xl"
     pinecone_index = pinecone.Index(pinecone_index_name)
-    retriever_model = get_instructor_embedding_model()
+    retriever_model = get_instructor_embedding_model_api()
     (
         sparse_retriever_model,
         sparse_retriever_tokenizer,
@@ -382,6 +400,7 @@ if document_type == "Single-Document":
         dense_query_embedding, sparse_query_embedding = hybrid_score_norm(
             dense_query_embedding, sparse_query_embedding, 0.3
         )
+
         query_results = query_pinecone_sparse(
             dense_query_embedding,
             sparse_query_embedding,
@@ -392,6 +411,7 @@ if document_type == "Single-Document":
             ticker,
             participant_type,
             keywords,
+            indices,
             threshold,
         )
 
@@ -413,6 +433,7 @@ if document_type == "Single-Document":
             ticker,
             participant_type,
             keywords,
+            indices,
             threshold,
         )
 
@@ -459,6 +480,7 @@ else:
                     ticker,
                     participant_type,
                     keywords,
+                    indices,
                     threshold,
                 )
                 results_list = sentence_id_combine(
@@ -490,6 +512,7 @@ else:
                     ticker,
                     participant_type,
                     keywords,
+                    indices,
                     threshold,
                 )
                 results_list = sentence_id_combine(
@@ -535,6 +558,7 @@ else:
                     ticker_first,
                     participant_type,
                     keywords,
+                    indices,
                     threshold,
                 )
                 results_list = sentence_id_combine(
@@ -557,6 +581,7 @@ else:
                     ticker_second,
                     participant_type,
                     keywords,
+                    indices,
                     threshold,
                 )
                 results_list = sentence_id_combine(
@@ -591,6 +616,7 @@ else:
                     ticker_first,
                     participant_type,
                     keywords,
+                    indices,
                     threshold,
                 )
                 results_list = sentence_id_combine(
@@ -612,6 +638,7 @@ else:
                     ticker_second,
                     participant_type,
                     keywords,
+                    indices,
                     threshold,
                 )
                 results_list = sentence_id_combine(
@@ -778,7 +805,7 @@ if decoder_model == "GPT-J":
             )
             submitted = st.form_submit_button("Submit")
 
-tab1, tab2 = st.tabs(["Retrived Text", "Retrieved Documents"])
+tab1, tab2 = st.tabs(["Retrieved Text", "Retrieved Documents"])
 
 
 with tab1:

requirements.txt

@@ -14,4 +14,4 @@ streamlit-scrollable-textbox
 openai
 InstructorEmbedding
 gradio_client
-
+rank_bm25

utils/models.py

@@ -20,26 +20,59 @@ from transformers import (
     T5Tokenizer,
     pipeline,
 )
-
-import pinecone
+from rank_bm25 import BM25Okapi, BM25L, BM25Plus
+import numpy as np
+from nltk.tokenize import word_tokenize
+from nltk.corpus import stopwords
+from nltk.stem.porter import PorterStemmer
+import re
 import streamlit as st
 
 
-@st.experimental_singleton
+@st.cache_resource
 def get_data():
     data = pd.read_csv("earnings_calls_cleaned_metadata.csv")
     return data
 
 
+# Preprocessing for BM25
+
+
+def tokenizer(
+    string, reg="[a-zA-Z'-]+|[0-9]{1,}%|[0-9]{1,}\.[0-9]{1,}%|\d+\.\d+%}"
+):
+    regex = reg
+    string = string.replace("-", " ")
+    return " ".join(re.findall(regex, string))
+
+
+def preprocess_text(text):
+    # Convert to lowercase
+    text = text.lower()
+    # Tokenize the text
+    tokens = word_tokenize(text)
+    # Remove stop words
+    stop_words = set(stopwords.words("english"))
+    tokens = [token for token in tokens if token not in stop_words]
+    # Stem the tokens
+    porter_stemmer = PorterStemmer()
+    tokens = [porter_stemmer.stem(token) for token in tokens]
+    # Join the tokens back into a single string
+    preprocessed_text = " ".join(tokens)
+    preprocessed_text = tokenizer(preprocessed_text)
+
+    return preprocessed_text
+
+
 # Initialize Spacy Model
 
 
-@st.experimental_singleton
+@st.cache_resource
 def get_spacy_model():
     return spacy.load("en_core_web_trf")
 
 
-@st.experimental_singleton
+@st.cache_resource
 def get_flan_alpaca_xl_model():
     model = AutoModelForSeq2SeqLM.from_pretrained(
         "/home/user/app/models/flan-alpaca-xl/"
@@ -53,19 +86,19 @@ def get_flan_alpaca_xl_model():
 # Initialize models from HuggingFace
 
 
-@st.experimental_singleton
+@st.cache_resource
 def get_t5_model():
     return pipeline("summarization", model="t5-small", tokenizer="t5-small")
 
 
-@st.experimental_singleton
+@st.cache_resource
 def get_flan_t5_model():
     tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
     model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large")
     return model, tokenizer
 
 
-@st.experimental_singleton
+@st.cache_resource
 def get_mpnet_embedding_model():
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = SentenceTransformer(
@@ -75,7 +108,7 @@ def get_mpnet_embedding_model():
     return model
 
 
-@st.experimental_singleton
+@st.cache_resource
 def get_splade_sparse_embedding_model():
     model_sparse = "naver/splade-cocondenser-ensembledistil"
     # check device
@@ -87,7 +120,7 @@ def get_splade_sparse_embedding_model():
     return model_sparse, tokenizer
 
 
-@st.experimental_singleton
+@st.cache_resource
 def get_sgpt_embedding_model():
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = SentenceTransformer(
@@ -97,20 +130,34 @@ def get_sgpt_embedding_model():
     return model
 
 
-@st.experimental_singleton
+@st.cache_resource
 def get_instructor_embedding_model():
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = INSTRUCTOR("hkunlp/instructor-xl")
     return model
 
+@st.cache_resource
+def get_instructor_embedding_model_api():
+    client = Client("https://awinml-api-instructor-xl-2.hf.space/")
+    return client
+
 
-@st.experimental_singleton
+@st.cache_resource
 def get_alpaca_model():
     client = Client("https://awinml-alpaca-cpp.hf.space")
     return client
 
 
-@st.experimental_memo
+@st.cache_resource
+def get_bm25_model(data):
+    corpus = data.Text.tolist()
+    corpus_clean = [preprocess_text(x) for x in corpus]
+    tokenized_corpus = [doc.split(" ") for doc in corpus_clean]
+    bm25 = BM25Plus(tokenized_corpus)
+    return corpus, bm25
+
+
+@st.cache_resource
 def save_key(api_key):
     return api_key
 

utils/nltkmodules.py

@@ -1,4 +1,5 @@
 import nltk
 
-nltk.download('wordnet')
-nltk.download('punkt')
+nltk.download("wordnet")
+nltk.download("punkt")
+nltk.download("stopwords")

utils/retriever.py

@@ -1,6 +1,16 @@
-def query_pinecone_sparse(
+def get_bm25_search_hits(corpus, sparse_scores, top_n=50):
+    bm25_search = []
+    indices = []
+    for idx in sparse_scores:
+        if len(bm25_search) <= top_n:
+            bm25_search.append(corpus[idx])
+            indices.append(idx)
+    indices = [int(x) for x in indices]
+    return indices
+
+
+def query_pinecone(
     dense_vec,
-    sparse_vec,
     top_k,
     index,
     year,
@@ -8,6 +18,7 @@ def query_pinecone_sparse(
     ticker,
     participant_type,
     keywords=None,
+    indices=None,
     threshold=0.25,
 ):
     if participant_type == "Company Speaker":
@@ -16,68 +27,126 @@ def query_pinecone_sparse(
         participant = "Question"
 
     # Create filter dictionary based on keywords
-    filter_dict = [{'Keywords': word} for word in keywords]
+    filter_dict = [{"Keywords": word} for word in keywords]
 
     if year == "All":
         if quarter == "All":
+            if indices != None:
+                xc = index.query(
+                    vector=dense_vec,
+                    top_k=top_k,
+                    filter={
+                        "Year": {
+                            "$in": [
+                                int("2020"),
+                                int("2019"),
+                                int("2018"),
+                                int("2017"),
+                                int("2016"),
+                            ]
+                        },
+                        "Quarter": {"$in": ["Q1", "Q2", "Q3", "Q4"]},
+                        "Ticker": {"$eq": ticker},
+                        "QA_Flag": {"$eq": participant},
+                        "Keywords": {"$in": keywords},
+                        "index": {"$in": indices},
+                    },
+                    include_metadata=True,
+                )
+            else:
+                xc = index.query(
+                    vector=dense_vec,
+                    top_k=top_k,
+                    filter={
+                        "Year": {
+                            "$in": [
+                                int("2020"),
+                                int("2019"),
+                                int("2018"),
+                                int("2017"),
+                                int("2016"),
+                            ]
+                        },
+                        "Quarter": {"$in": ["Q1", "Q2", "Q3", "Q4"]},
+                        "Ticker": {"$eq": ticker},
+                        "QA_Flag": {"$eq": participant},
+                        "Keywords": {"$in": keywords},
+                    },
+                    include_metadata=True,
+                )
+        else:
+            if indices != None:
+                xc = index.query(
+                    vector=dense_vec,
+                    top_k=top_k,
+                    filter={
+                        "Year": {
+                            "$in": [
+                                int("2020"),
+                                int("2019"),
+                                int("2018"),
+                                int("2017"),
+                                int("2016"),
+                            ]
+                        },
+                        "Quarter": {"$eq": quarter},
+                        "Ticker": {"$eq": ticker},
+                        "QA_Flag": {"$eq": participant},
+                        "Keywords": {"$in": keywords},
+                        "index": {"$in": indices},
+                    },
+                    include_metadata=True,
+                )
+            else:
+                xc = index.query(
+                    vector=dense_vec,
+                    top_k=top_k,
+                    filter={
+                        "Year": {
+                            "$in": [
+                                int("2020"),
+                                int("2019"),
+                                int("2018"),
+                                int("2017"),
+                                int("2016"),
+                            ]
+                        },
+                        "Quarter": {"$eq": quarter},
+                        "Ticker": {"$eq": ticker},
+                        "QA_Flag": {"$eq": participant},
+                        "Keywords": {"$in": keywords},
+                    },
+                    include_metadata=True,
+                )
+    else:
+        # search pinecone index for context passage with the answer
+        if indices != None:
             xc = index.query(
                 vector=dense_vec,
-                sparse_vector=sparse_vec,
                 top_k=top_k,
                 filter={
-                    "Year": {
-                        "$in": [
-                            int("2020"),
-                            int("2019"),
-                            int("2018"),
-                            int("2017"),
-                            int("2016"),
-                        ]
-                    },
-                    "Quarter": {"$in": ["Q1", "Q2", "Q3", "Q4"]},
+                    "Year": int(year),
+                    "Quarter": {"$eq": quarter},
                     "Ticker": {"$eq": ticker},
                     "QA_Flag": {"$eq": participant},
-                    "Keywords": {"$in": keywords}
+                    "Keywords": {"$in": keywords},
+                    "index": {"$in": indices},
                 },
                 include_metadata=True,
             )
         else:
             xc = index.query(
                 vector=dense_vec,
-                sparse_vector=sparse_vec,
                 top_k=top_k,
                 filter={
-                    "Year": {
-                        "$in": [
-                            int("2020"),
-                            int("2019"),
-                            int("2018"),
-                            int("2017"),
-                            int("2016"),
-                        ]
-                    },
+                    "Year": int(year),
                     "Quarter": {"$eq": quarter},
                     "Ticker": {"$eq": ticker},
                     "QA_Flag": {"$eq": participant},
-                    "Keywords": {"$in": keywords}
+                    "Keywords": {"$in": keywords},
                 },
                 include_metadata=True,
             )
-    else:
-        # search pinecone index for context passage with the answer
-        xc = index.query(
-            vector=dense_vec,
-            sparse_vector=sparse_vec,
-            top_k=top_k,
-            filter={
-                "Year": int(year),
-                "Quarter": {"$eq": quarter},
-                "Ticker": {"$eq": ticker},
-                "QA_Flag": {"$eq": participant},
-                "Keywords": {"$in": keywords}
-            },
-            include_metadata=True,
-        )
     # filter the context passages based on the score threshold
     filtered_matches = []
     for match in xc["matches"]:
@@ -87,8 +156,9 @@ def query_pinecone_sparse(
     return xc
 
 
-def query_pinecone(
+def query_pinecone_sparse(
     dense_vec,
+    sparse_vec,
     top_k,
     index,
     year,
@@ -96,6 +166,7 @@ def query_pinecone(
     ticker,
     participant_type,
     keywords=None,
+    indices=None,
     threshold=0.25,
 ):
     if participant_type == "Company Speaker":
@@ -104,13 +175,13 @@ def query_pinecone(
         participant = "Question"
 
     # Create filter dictionary based on keywords
-    filter_dict = [{'Keywords': word} for word in keywords]
-
+    filter_dict = [{"Keywords": word} for word in keywords]
 
     if year == "All":
         if quarter == "All":
             xc = index.query(
                 vector=dense_vec,
+                sparse_vector=sparse_vec,
                 top_k=top_k,
                 filter={
                     "Year": {
@@ -125,13 +196,14 @@ def query_pinecone(
                     "Quarter": {"$in": ["Q1", "Q2", "Q3", "Q4"]},
                     "Ticker": {"$eq": ticker},
                     "QA_Flag": {"$eq": participant},
-                    "Keywords": {"$in": keywords}
+                    "Keywords": {"$in": keywords},
                 },
                 include_metadata=True,
             )
         else:
             xc = index.query(
                 vector=dense_vec,
+                sparse_vector=sparse_vec,
                 top_k=top_k,
                 filter={
                     "Year": {
@@ -146,7 +218,7 @@ def query_pinecone(
                     "Quarter": {"$eq": quarter},
                     "Ticker": {"$eq": ticker},
                     "QA_Flag": {"$eq": participant},
-                    "Keywords": {"$in": keywords}
+                    "Keywords": {"$in": keywords},
                 },
                 include_metadata=True,
             )
@@ -154,13 +226,14 @@ def query_pinecone(
         # search pinecone index for context passage with the answer
         xc = index.query(
             vector=dense_vec,
+            sparse_vector=sparse_vec,
             top_k=top_k,
             filter={
                 "Year": int(year),
                 "Quarter": {"$eq": quarter},
                 "Ticker": {"$eq": ticker},
                 "QA_Flag": {"$eq": participant},
-                "Keywords": {"$in": keywords}
+                "Keywords": {"$in": keywords},
             },
             include_metadata=True,
         )

utils/vector_index.py

@@ -1,11 +1,23 @@
 import torch
+import json
+import numpy as np
 
 
 def create_dense_embeddings(query, model, instruction=None):
     if instruction == None:
         dense_emb = model.encode([query]).tolist()
     else:
-        dense_emb = model.encode([[instruction, query]]).tolist()
+        # Fetching embedding from API for Instructor
+        json_output_embedding = model.predict(
+            instruction,
+            query,
+            api_name="/predict",
+        )
+
+        json_file = open(json_output_embedding, "r")
+        json_dict = json.load(json_file)
+        dense_array = np.array(json_dict["data"], dtype=np.float64)
+        dense_emb = dense_array.tolist()
     return dense_emb