Upload 8 files

app.py

@@ -1,3 +1,5 @@
+import re
+
 import openai
 import streamlit_scrollable_textbox as stx
 
@@ -25,7 +27,7 @@ from utils.models import (
     get_spacy_model,
     get_splade_sparse_embedding_model,
     get_t5_model,
-    gpt_model,
+    gpt_turbo_model,
     save_key,
 )
 from utils.prompts import (
@@ -36,8 +38,10 @@ from utils.prompts import (
     generate_flant5_prompt_summ_chunk_context_single,
     generate_gpt_j_two_shot_prompt_1,
     generate_gpt_j_two_shot_prompt_2,
-    generate_gpt_prompt,
-    generate_gpt_prompt_2,
+    generate_gpt_prompt_alpaca,
+    generate_gpt_prompt_alpaca_multi_doc,
+    generate_gpt_prompt_original,
+    generate_multi_doc_context,
     get_context_list_prompt,
 )
 from utils.retriever import (
@@ -46,6 +50,7 @@ from utils.retriever import (
     query_pinecone_sparse,
     sentence_id_combine,
     text_lookup,
+    year_quarter_range,
 )
 from utils.transcript_retrieval import retrieve_transcript
 from utils.vector_index import (
@@ -66,59 +71,29 @@ col1, col2 = st.columns([3, 3], gap="medium")
 
 with st.sidebar:
     ner_choice = st.selectbox("Select NER Model", ["Spacy", "Alpaca"])
+    document_type = st.selectbox(
+        "Select Query Type", ["Single-Document", "Multi-Document"]
+    )
 
 if ner_choice == "Spacy":
     ner_model = get_spacy_model()
 
 with col1:
     st.subheader("Question")
-    query_text = st.text_area(
-        "Input Query",
-        value="What was discussed regarding Wearables revenue performance?",
-    )
-
-if ner_choice == "Alpaca":
-    ner_prompt = generate_alpaca_ner_prompt(query_text)
-    entity_text = generate_entities_flan_alpaca_inference_api(ner_prompt)
-    company_ent, quarter_ent, year_ent = format_entities_flan_alpaca(
-        entity_text
-    )
-else:
-    company_ent = extract_ticker_spacy(query_text, ner_model)
-    quarter_ent, year_ent = extract_quarter_year(query_text)
-
-ticker_index, quarter_index, year_index = clean_entities(
-    company_ent, quarter_ent, year_ent
-)
-
-with col1:
-    years_choice = ["2020", "2019", "2018", "2017", "2016", "All"]
-
-with col1:
-    # Hardcoding the defaults for a question without metadata
-    if (
-        query_text
-        == "What was discussed regarding Wearables revenue performance?"
-    ):
-        year = st.selectbox("Year", years_choice)
-    else:
-        year = st.selectbox("Year", years_choice, index=year_index)
-
-with col1:
-    # Hardcoding the defaults for a question without metadata
-    if (
-        query_text
-        == "What was discussed regarding Wearables revenue performance?"
-    ):
-        quarter = st.selectbox("Quarter", ["Q1", "Q2", "Q3", "Q4", "All"])
+    if document_type == "Single-Document":
+        query_text = st.text_area(
+            "Input Query",
+            value="What was discussed regarding Wearables revenue performance?",
+        )
     else:
-        quarter = st.selectbox(
-            "Quarter", ["Q1", "Q2", "Q3", "Q4", "All"], index=quarter_index
+        query_text = st.text_area(
+            "Input Query",
+            value="How has Apple's revenue from Wearables performed over the past 2 years?",
         )
 
-with col1:
-    participant_type = st.selectbox("Speaker", ["Company Speaker", "Analyst"])
 
+years_choice = ["2020", "2019", "2018", "2017", "2016", "All"]
+quarters_choice = ["Q1", "Q2", "Q3", "Q4", "All"]
 ticker_choice = [
     "AAPL",
     "CSCO",
@@ -132,23 +107,87 @@ ticker_choice = [
     "AMD",
 ]
 
-with col1:
-    # Hardcoding the defaults for a question without metadata
-    if (
-        query_text
-        == "What was discussed regarding Wearables revenue performance?"
-    ):
-        ticker = st.selectbox("Company", ticker_choice)
+
+if document_type == "Single-Document":
+    if ner_choice == "Alpaca":
+        ner_prompt = generate_alpaca_ner_prompt(query_text)
+        entity_text = generate_entities_flan_alpaca_inference_api(ner_prompt)
+        company_ent, quarter_ent, year_ent = format_entities_flan_alpaca(
+            entity_text
+        )
     else:
-        ticker = st.selectbox("Company", ticker_choice, ticker_index)
+        company_ent = extract_ticker_spacy(query_text, ner_model)
+        quarter_ent, year_ent = extract_quarter_year(query_text)
+
+    ticker_index, quarter_index, year_index = clean_entities(
+        company_ent, quarter_ent, year_ent
+    )
+
+    with col1:
+        # Hardcoding the defaults for a question without metadata
+        if (
+            query_text
+            == "What was discussed regarding Wearables revenue performance?"
+        ):
+            year = st.selectbox("Year", years_choice)
+            quarter = st.selectbox("Quarter", quarters_choice)
+            ticker = st.selectbox("Company", ticker_choice)
+        else:
+            year = st.selectbox("Year", years_choice, index=year_index)
+            quarter = st.selectbox(
+                "Quarter", quarters_choice, index=quarter_index
+            )
+            ticker = st.selectbox("Company", ticker_choice, ticker_index)
+
+        participant_type = st.selectbox(
+            "Speaker", ["Company Speaker", "Analyst"]
+        )
+
+else:
+    # Multi-Document Case
+
+    with col1:
+        # Hardcoding the defaults for a question without metadata
+        if (
+            query_text
+            == "How has Apple's revenue from Wearables performed over the past 2 years?"
+        ):
+            start_year = st.selectbox("Start Year", years_choice, index=2)
+            start_quarter = st.selectbox(
+                "Start Quarter", quarters_choice, index=0
+            )
+
+            end_year = st.selectbox("End Year", years_choice, index=0)
+            end_quarter = st.selectbox("End Quarter", quarters_choice, index=0)
+
+            ticker = st.selectbox("Company", ticker_choice, index=0)
+        else:
+            start_year = st.selectbox("Start Year", years_choice, index=2)
+            start_quarter = st.selectbox(
+                "Start Quarter", quarters_choice, index=0
+            )
+
+            end_year = st.selectbox("End Year", years_choice, index=0)
+            end_quarter = st.selectbox("End Quarter", quarters_choice, index=0)
+
+            ticker = st.selectbox("Company", ticker_choice, index=0)
+
+        participant_type = st.selectbox(
+            "Speaker", ["Company Speaker", "Analyst"]
+        )
+
 
 with st.sidebar:
     st.subheader("Select Options:")
 
-with st.sidebar:
-    num_results = int(
-        st.number_input("Number of Results to query", 1, 15, value=5)
-    )
+    if document_type == "Single-Document":
+        num_results = int(
+            st.number_input("Number of Results to query", 1, 15, value=5)
+        )
+    else:
+        num_results = int(
+            st.number_input("Number of Results to query", 1, 15, value=2)
+        )
 
 
 # Choose encoder model
@@ -160,8 +199,11 @@ with st.sidebar:
 
 # Choose decoder model
 
-decoder_models_choice = ["GPT3 - (text-davinci-003)", "T5", "FLAN-T5", "GPT-J"]
-
+# Restricting multi-document to only GPT-3
+if document_type == "Single-Document":
+    decoder_models_choice = ["GPT-3.5 Turbo", "T5", "FLAN-T5", "GPT-J"]
+else:
+    decoder_models_choice = ["GPT-3.5 Turbo"]
 with st.sidebar:
     decoder_model = st.selectbox("Select Decoder Model", decoder_models_choice)
 
@@ -198,66 +240,140 @@ elif encoder_model == "Hybrid MPNET - SPLADE":
     ) = get_splade_sparse_embedding_model()
 
 with st.sidebar:
-    window = int(st.number_input("Sentence Window Size", 0, 10, value=1))
-
-with st.sidebar:
-    threshold = float(
-        st.number_input(
-            label="Similarity Score Threshold",
-            step=0.05,
-            format="%.2f",
-            value=0.25,
+    if document_type == "Single-Document":
+        window = int(st.number_input("Sentence Window Size", 0, 10, value=1))
+
+        threshold = float(
+            st.number_input(
+                label="Similarity Score Threshold",
+                step=0.05,
+                format="%.2f",
+                value=0.25,
+            )
+        )
+    else:
+        window = int(st.number_input("Sentence Window Size", 0, 10, value=0))
+
+        threshold = float(
+            st.number_input(
+                label="Similarity Score Threshold",
+                step=0.05,
+                format="%.2f",
+                value=0.6,
+            )
         )
-    )
 
 data = get_data()
 
-if encoder_model == "Hybrid SGPT - SPLADE":
-    dense_query_embedding = create_dense_embeddings(
-        query_text, retriever_model
-    )
-    sparse_query_embedding = create_sparse_embeddings(
-        query_text, sparse_retriever_model, sparse_retriever_tokenizer
-    )
-    dense_query_embedding, sparse_query_embedding = hybrid_score_norm(
-        dense_query_embedding, sparse_query_embedding, 0
-    )
-    query_results = query_pinecone_sparse(
-        dense_query_embedding,
-        sparse_query_embedding,
-        num_results,
-        pinecone_index,
-        year,
-        quarter,
-        ticker,
-        participant_type,
-        threshold,
-    )
+if document_type == "Single-Document":
+    if encoder_model == "Hybrid SGPT - SPLADE":
+        dense_query_embedding = create_dense_embeddings(
+            query_text, retriever_model
+        )
+        sparse_query_embedding = create_sparse_embeddings(
+            query_text, sparse_retriever_model, sparse_retriever_tokenizer
+        )
+        dense_query_embedding, sparse_query_embedding = hybrid_score_norm(
+            dense_query_embedding, sparse_query_embedding, 0
+        )
+        query_results = query_pinecone_sparse(
+            dense_query_embedding,
+            sparse_query_embedding,
+            num_results,
+            pinecone_index,
+            year,
+            quarter,
+            ticker,
+            participant_type,
+            threshold,
+        )
 
-else:
-    dense_query_embedding = create_dense_embeddings(
-        query_text, retriever_model
-    )
-    query_results = query_pinecone(
-        dense_query_embedding,
-        num_results,
-        pinecone_index,
-        year,
-        quarter,
-        ticker,
-        participant_type,
-        threshold,
-    )
+    else:
+        dense_query_embedding = create_dense_embeddings(
+            query_text, retriever_model
+        )
+        query_results = query_pinecone(
+            dense_query_embedding,
+            num_results,
+            pinecone_index,
+            year,
+            quarter,
+            ticker,
+            participant_type,
+            threshold,
+        )
 
+    if threshold <= 0.90:
+        context_list = sentence_id_combine(data, query_results, lag=window)
+    else:
+        context_list = format_query(query_results)
 
-if threshold <= 0.90:
-    context_list = sentence_id_combine(data, query_results, lag=window)
 else:
-    context_list = format_query(query_results)
+    # Multi-Document Retreival
+    if encoder_model == "Hybrid SGPT - SPLADE":
+        dense_query_embedding = create_dense_embeddings(
+            query_text, retriever_model
+        )
+        sparse_query_embedding = create_sparse_embeddings(
+            query_text, sparse_retriever_model, sparse_retriever_tokenizer
+        )
+        dense_query_embedding, sparse_query_embedding = hybrid_score_norm(
+            dense_query_embedding, sparse_query_embedding, 0
+        )
+        year_quarter_list = year_quarter_range(
+            start_quarter, start_year, end_quarter, end_year
+        )
+
+        context_group = []
+        for year, quarter in year_quarter_list:
+            query_results = query_pinecone_sparse(
+                dense_query_embedding,
+                sparse_query_embedding,
+                num_results,
+                pinecone_index,
+                year,
+                quarter,
+                ticker,
+                participant_type,
+                threshold,
+            )
+            results_list = sentence_id_combine(data, query_results, lag=window)
+            context_group.append((results_list, year, quarter))
+
+    else:
+        dense_query_embedding = create_dense_embeddings(
+            query_text, retriever_model
+        )
+        year_quarter_list = year_quarter_range(
+            start_quarter, start_year, end_quarter, end_year
+        )
+
+        context_group = []
+        for year, quarter in year_quarter_list:
+            query_results = query_pinecone(
+                dense_query_embedding,
+                num_results,
+                pinecone_index,
+                year,
+                quarter,
+                ticker,
+                participant_type,
+                threshold,
+            )
+            results_list = sentence_id_combine(data, query_results, lag=window)
+            context_group.append((results_list, year, quarter))
 
+    multi_doc_context = generate_multi_doc_context(context_group)
+
+
+if decoder_model == "GPT-3.5 Turbo":
+    if document_type == "Single-Document":
+        prompt = generate_gpt_prompt_alpaca(query_text, context_list)
+    else:
+        prompt = generate_gpt_prompt_alpaca_multi_doc(
+            query_text, context_group
+        )
 
-if decoder_model == "GPT3 - (text-davinci-003)":
-    prompt = generate_gpt_prompt(query_text, context_list)
     with col2:
         with st.form("my_form"):
             edited_prompt = st.text_area(
@@ -273,9 +389,20 @@ if decoder_model == "GPT3 - (text-davinci-003)":
             if submitted:
                 api_key = save_key(openai_key)
                 openai.api_key = api_key
-                generated_text = gpt_model(edited_prompt)
+                generated_text = gpt_turbo_model(edited_prompt)
                 st.subheader("Answer:")
-                st.write(generated_text)
+                regex_pattern_sentences = (
+                    "(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?)\s"
+                )
+                generated_text_list = re.split(
+                    regex_pattern_sentences, generated_text
+                )
+                for answer_text in generated_text_list:
+                    answer_text = f"""{answer_text}"""
+                    st.write(
+                        f"<ul><li><p>{answer_text}</p></li></ul>",
+                        unsafe_allow_html=True,
+                    )
 
 
 elif decoder_model == "T5":
@@ -384,22 +511,53 @@ if decoder_model == "GPT-J":
             )
             submitted = st.form_submit_button("Submit")
 
+tab1, tab2 = st.tabs(["Retrived Text", "Retrieved Documents"])
 
-with col1:
-    with st.expander("See Retrieved Text"):
-        st.subheader("Retrieved Text:")
-        for context_text in context_list:
-            context_text = f"""{context_text}"""
-            st.write(
-                f"<ul><li><p>{context_text}</p></li></ul>",
-                unsafe_allow_html=True,
-            )
-
-file_text = retrieve_transcript(data, year, quarter, ticker)
 
-with col1:
-    with st.expander("See Transcript"):
-        st.subheader("Earnings Call Transcript:")
-        stx.scrollableTextbox(
-            file_text, height=700, border=False, fontFamily="Helvetica"
-        )
+with tab1:
+    if document_type == "Single-Document":
+        with st.expander("See Retrieved Text"):
+            st.subheader("Retrieved Text:")
+            for context_text in context_list:
+                context_text = f"""{context_text}"""
+                st.write(
+                    f"<ul><li><p>{context_text}</p></li></ul>",
+                    unsafe_allow_html=True,
+                )
+    else:
+        with st.expander("See Retrieved Text"):
+            st.subheader("Retrieved Text:")
+            sections = [
+                s.strip()
+                for s in multi_doc_context.split("Document: ")
+                if s.strip()
+            ]
+
+            # Add "Document: " back to the beginning of each section
+            context_list = [
+                "Document: " + s[0:7] + "\n" + s[7:] for s in sections
+            ]
+            for context_text in context_list:
+                context_text = f"""{context_text}"""
+                st.write(
+                    f"<ul><li><p>{context_text}</p></li></ul>",
+                    unsafe_allow_html=True,
+                )
+
+
+with tab2:
+    if document_type == "Single-Document":
+        file_text = retrieve_transcript(data, year, quarter, ticker)
+        with st.expander("See Transcript"):
+            st.subheader("Earnings Call Transcript:")
+            stx.scrollableTextbox(
+                file_text, height=700, border=False, fontFamily="Helvetica"
+            )
+    else:
+        for year, quarter in year_quarter_list:
+            file_text = retrieve_transcript(data, year, quarter, ticker)
+            with st.expander(f"See Transcript - {quarter} {year}"):
+                st.subheader("Earnings Call Transcript - {quarter} {year}:")
+                stx.scrollableTextbox(
+                    file_text, height=700, border=False, fontFamily="Helvetica"
+                )

utils/models.py

@@ -103,14 +103,16 @@ def save_key(api_key):
 # Text Generation
 
 
-def gpt_model(prompt):
-    response = openai.Completion.create(
-        model="text-davinci-003",
-        prompt=prompt,
-        temperature=0,
+def gpt_turbo_model(prompt):
+    response = openai.ChatCompletion.create(
+        model="gpt-3.5-turbo",
+        messages=[
+            {"role": "user", "content": prompt},
+        ],
+        temperature=0.01,
         max_tokens=1024,
     )
-    return response.choices[0].text
+    return response["choices"][0]["message"]["content"]
 
 
 def generate_text_flan_t5(model, tokenizer, input_text):

utils/prompts.py

@@ -1,4 +1,51 @@
-def generate_gpt_prompt(query_text, context_list):
+def generate_multi_doc_context(context_group):
+    multi_doc_context = ""
+    for context_text_list, year, quarter in context_group:
+        print((context_text_list, year, quarter))
+        if context_text_list == []:
+            break
+        else:
+            multi_doc_context = (
+                multi_doc_context
+                + "\n"
+                + f"Document: {quarter} {year}"
+                + "\n"
+                + " ".join(context_text_list)
+            )
+    return multi_doc_context
+
+
+def generate_gpt_prompt_alpaca(query_text, context_list):
+    context = " ".join(context_list)
+    prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Use the following guidelines to write a response that that appropriately completes the request:
+### Instruction:
+- Write a detailed paragraph consisting of exactly five complete sentences that answer the question based on the provided context.
+- Focus on addressing the specific question posed, providing as much relevant information and detail as possible.
+- Only use details from the provided context that directly address the question; do not include any additional information that is not explicitly stated.
+- Aim to provide a clear and concise summary that fully addresses the question.
+
+Question: {query_text}
+Context: {context}
+### Response:"""
+    return prompt
+
+
+def generate_gpt_prompt_alpaca_multi_doc(query_text, context_group):
+    multi_doc_context = generate_multi_doc_context(context_group)
+    prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Use the following guidelines to write a response that that appropriately completes the request:
+### Instruction:
+- Write a detailed paragraph consisting of exactly five complete sentences that answer the question based on the provided context.
+- Focus on addressing the specific question posed, providing as much relevant information and detail as possible.
+- Only use details from the provided context that directly address the question; do not include any additional information that is not explicitly stated.
+- Aim to provide a clear and concise summary that fully addresses the question.
+
+Question: {query_text}
+Context: {multi_doc_context}
+### Response:"""
+    return prompt
+
+
+def generate_gpt_prompt_original(query_text, context_list):
     context = " ".join(context_list)
     prompt = f"""Answer the question in 6 long detailed points as accurately as possible using the provided context. Include as many key details as possible.
 Context: {context}

utils/retriever.py

@@ -195,3 +195,55 @@ def sentence_id_combine(data, query_results, lag=1):
 def text_lookup(data, sentence_ids):
     context = ". ".join(data.iloc[sentence_ids].to_list())
     return context
+
+
+def year_quarter_range(start_quarter, start_year, end_quarter, end_year):
+    """Creates a list of all (year, quarter) pairs that lie in the range including the start and end quarters."""
+    start_year = int(start_year)
+    end_year = int(end_year)
+
+    quarters = (
+        [("Q1", "Q2", "Q3", "Q4")] * (end_year - start_year)
+        + [("Q1", "Q2", "Q3" if end_quarter == "Q4" else "Q4")]
+        * (end_quarter == "Q4")
+        + [
+            (
+                "Q1"
+                if start_quarter == "Q1"
+                else "Q2"
+                if start_quarter == "Q2"
+                else "Q3"
+                if start_quarter == "Q3"
+                else "Q4",
+            )
+            * (end_year - start_year)
+        ]
+    )
+    years = list(range(start_year, end_year + 1))
+    list_year_quarter = [
+        (y, q) for y in years for q in quarters[years.index(y)]
+    ]
+    # Remove duplicate pairs
+    seen = set()
+    list_year_quarter_cleaned = []
+    for tup in list_year_quarter:
+        if tup not in seen:
+            seen.add(tup)
+            list_year_quarter_cleaned.append(tup)
+    return list_year_quarter_cleaned
+
+
+def multi_document_query(
+    dense_query_embedding,
+    sparse_query_embedding,
+    num_results,
+    pinecone_index,
+    start_quarter,
+    start_year,
+    end_quarter,
+    end_year,
+    ticker,
+    participant_type,
+    threshold,
+):
+    pass