Upload app.py

app.py

@@ -40,6 +40,7 @@ from utils.prompts import (
     generate_gpt_j_two_shot_prompt_2,
     generate_gpt_prompt_alpaca,
     generate_gpt_prompt_alpaca_multi_doc,
+    generate_gpt_prompt_alpaca_multi_doc_multi_company,
     generate_gpt_prompt_original,
     generate_multi_doc_context,
     get_context_list_prompt,
@@ -74,6 +75,11 @@ with st.sidebar:
     document_type = st.selectbox(
         "Select Query Type", ["Single-Document", "Multi-Document"]
     )
+    if document_type == "Multi-Document":
+        multi_company_choice = st.selectbox(
+            "Select Company Query Type",
+            ["Single-Company", "Compare Companies"],
+        )
 
 if ner_choice == "Spacy":
     ner_model = get_spacy_model()
@@ -86,11 +92,16 @@ with col1:
             value="What was discussed regarding Wearables revenue performance?",
         )
     else:
-        query_text = st.text_area(
-            "Input Query",
-            value="How has revenue from Wearables performed over the past 2 years?",
-        )
-
+        if multi_company_choice == "Single-Company":
+            query_text = st.text_area(
+                "Input Query",
+                value="What was the reported revenue for Wearables over the last 2 years?",
+            )
+        else:
+            query_text = st.text_area(
+                "Input Query",
+                value="How was AAPL's capex spend compared to GOOGL?",
+            )
 
 years_choice = ["2020", "2019", "2018", "2017", "2016", "All"]
 quarters_choice = ["Q1", "Q2", "Q3", "Q4", "All"]
@@ -145,32 +156,76 @@ if document_type == "Single-Document":
 
 else:
     # Multi-Document Case
-
     with col1:
-        # Hardcoding the defaults for a question without metadata
-        if (
-            query_text
-            == "How has 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
-            )
+        # Single Company Summary
+        if multi_company_choice == "Single-Company":
+            # Hardcoding the defaults for a question without metadata
+            if (
+                query_text
+                == "What was the reported revenue for Wearables over the last 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)
+                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
-            )
+                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)
+
+        # Single Company Summary
+        if multi_company_choice == "Compare Companies":
+            # Hardcoding the defaults for a question without metadata
+            if query_text == "How was AAPL's capex spend compared to GOOGL?":
+                start_year = st.selectbox("Start Year", years_choice, index=1)
+                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)
+                end_year = st.selectbox("End Year", years_choice, index=0)
+                end_quarter = st.selectbox(
+                    "End Quarter", quarters_choice, index=0
+                )
+
+                ticker_first = st.selectbox(
+                    "First Company", ticker_choice, index=0
+                )
+                ticker_second = st.selectbox(
+                    "Second Company", ticker_choice, index=5
+                )
+
+            else:
+                start_year = st.selectbox("Start Year", years_choice, index=2)
+                start_quarter = st.selectbox(
+                    "Start Quarter", quarters_choice, index=0
+                )
 
-            ticker = st.selectbox("Company", ticker_choice, index=0)
+                end_year = st.selectbox("End Year", years_choice, index=0)
+                end_quarter = st.selectbox(
+                    "End Quarter", quarters_choice, index=0
+                )
+
+                ticker_first = st.selectbox(
+                    "First Company", ticker_choice, index=0
+                )
+                ticker_second = st.selectbox(
+                    "Second Company", ticker_choice, index=1
+                )
 
         participant_type = st.selectbox(
             "Speaker", ["Company Speaker", "Analyst"]
@@ -186,7 +241,7 @@ with st.sidebar:
         )
     else:
         num_results = int(
-            st.number_input("Number of Results to query", 1, 15, value=2)
+            st.number_input("Number of Results to query", 1, 15, value=4)
         )
 
 
@@ -252,7 +307,7 @@ with st.sidebar:
             )
         )
     else:
-        window = int(st.number_input("Sentence Window Size", 0, 10, value=0))
+        window = int(st.number_input("Sentence Window Size", 0, 10, value=1))
 
         threshold = float(
             st.number_input(
@@ -310,69 +365,191 @@ if document_type == "Single-Document":
 
 else:
     # 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
-        )
+    # Single Company
+    if multi_company_choice == "Single-Company":
+        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, ticker))
 
-        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,
+        else:
+            dense_query_embedding = create_dense_embeddings(
+                query_text, retriever_model
             )
-            results_list = sentence_id_combine(data, query_results, lag=window)
-            context_group.append((results_list, year, quarter))
+            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, ticker))
 
+        multi_doc_context = generate_multi_doc_context(context_group)
+    # Companies Comparison
     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
-        )
+        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
+            )
+
+            # First Company Context
+            context_group_first = []
+            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_first,
+                    participant_type,
+                    threshold,
+                )
+                results_list = sentence_id_combine(
+                    data, query_results, lag=window
+                )
+                context_group_first.append(
+                    (results_list, year, quarter, ticker_first)
+                )
+
+            # Second Company Context
+            context_group_second = []
+            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_second,
+                    participant_type,
+                    threshold,
+                )
+                results_list = sentence_id_combine(
+                    data, query_results, lag=window
+                )
+                context_group_second.append(
+                    (results_list, year, quarter, ticker_second)
+                )
 
-        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,
+        else:
+            dense_query_embedding = create_dense_embeddings(
+                query_text, retriever_model
             )
-            results_list = sentence_id_combine(data, query_results, lag=window)
-            context_group.append((results_list, year, quarter))
+            year_quarter_list = year_quarter_range(
+                start_quarter, start_year, end_quarter, end_year
+            )
+
+            # First Company Context
+            context_group_first = []
+            for year, quarter in year_quarter_list:
+                query_results = query_pinecone(
+                    dense_query_embedding,
+                    num_results,
+                    pinecone_index,
+                    year,
+                    quarter,
+                    ticker_first,
+                    participant_type,
+                    threshold,
+                )
+                results_list = sentence_id_combine(
+                    data, query_results, lag=window
+                )
+                context_group_first.append(
+                    (results_list, year, quarter, ticker_first)
+                )
 
-    multi_doc_context = generate_multi_doc_context(context_group)
+            # Second Company Context
+            context_group_second = []
+            for year, quarter in year_quarter_list:
+                query_results = query_pinecone(
+                    dense_query_embedding,
+                    num_results,
+                    pinecone_index,
+                    year,
+                    quarter,
+                    ticker_second,
+                    participant_type,
+                    threshold,
+                )
+                results_list = sentence_id_combine(
+                    data, query_results, lag=window
+                )
+                context_group_second.append(
+                    (results_list, year, quarter, ticker_second)
+                )
 
+        multi_doc_context_first = generate_multi_doc_context(
+            context_group_first
+        )
+        multi_doc_context_second = generate_multi_doc_context(
+            context_group_second
+        )
 
 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 multi_company_choice == "Single-Company":
+            prompt = generate_gpt_prompt_alpaca_multi_doc(
+                query_text, context_group
+            )
+        else:
+            prompt = generate_gpt_prompt_alpaca_multi_doc_multi_company(
+                query_text, context_group_first, context_group_second
+            )
 
     with col2:
         with st.form("my_form"):
@@ -527,6 +704,11 @@ with tab1:
     else:
         with st.expander("See Retrieved Text"):
             st.subheader("Retrieved Text:")
+            if multi_company_choice == "Compare Companies":
+                multi_doc_context = (
+                    multi_doc_context_first + multi_doc_context_second
+                )
+
             sections = [
                 s.strip()
                 for s in multi_doc_context.split("Document: ")
@@ -554,10 +736,26 @@ with tab2:
                 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"
-                )
+        if multi_company_choice == "Single-Company":
+            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"
+                    )
+        else:
+            for year, quarter in year_quarter_list:
+                file_text = retrieve_transcript(data, year, quarter, ticker_first)
+                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"
+                    )
+            for year, quarter in year_quarter_list:
+                file_text = retrieve_transcript(data, year, quarter, ticker_second)
+                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"
+                    )