Upload 16 files (#11)

- Upload 16 files (75295ea33fd86490c00779398a7101ad71ee52e9)

app.py

@@ -14,12 +14,14 @@ from utils.entity_extraction import (
     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_flan_alpaca_xl_model,
     get_flan_t5_model,
     get_instructor_embedding_model,
@@ -85,6 +87,8 @@ with st.sidebar:
 if ner_choice == "Spacy":
     ner_model = get_spacy_model()
 
+alpaca_model = get_alpaca_model()
+
 with col1:
     st.subheader("Question")
     if document_type == "Single-Document":
@@ -104,6 +108,10 @@ with col1:
                 value="How was AAPL's capex spend compared to GOOGL?",
             )
 
+
+# Extract keywords from query
+keywords = extract_keywords(query_text, alpaca_model)
+
 years_choice = ["2020", "2019", "2018", "2017", "2016", "All"]
 quarters_choice = ["Q1", "Q2", "Q3", "Q4", "All"]
 ticker_choice = [
@@ -382,6 +390,7 @@ if document_type == "Single-Document":
             quarter,
             ticker,
             participant_type,
+            keywords,
             threshold,
         )
 

requirements.txt

@@ -1,4 +1,5 @@
 pandas
+nltk
 tqdm
 pinecone-client
 spacy[transformers] == 3.3.0
@@ -12,3 +13,5 @@ streamlit
 streamlit-scrollable-textbox
 openai
 InstructorEmbedding
+gradio_client
+

utils/entity_extraction.py

@@ -1,4 +1,54 @@
 import re
+from nltk.stem import PorterStemmer, WordNetLemmatizer
+
+# Keyword Extracttion
+
+
+def expand_list_of_lists(list_of_lists):
+    """
+    Expands a list of lists of strings to a list of strings.
+    Args:
+      list_of_lists: A list of lists of strings.
+    Returns:
+      A list of strings.
+    """
+
+    expanded_list = []
+    for inner_list in list_of_lists:
+        for string in inner_list:
+            expanded_list.append(string)
+    return expanded_list
+
+
+def all_keywords_combs(texts):
+
+    texts = [text.split(" ") for text in texts]
+    texts = expand_list_of_lists(texts)
+
+    # Convert all strings to lowercase.
+    lower_texts = [text.lower() for text in texts]
+
+    # Stem the words in each string.
+    stemmer = PorterStemmer()
+    stem_texts = [stemmer.stem(text) for text in texts]
+
+    # Lemmatize the words in each string.
+    lemmatizer = WordNetLemmatizer()
+    lemm_texts = [lemmatizer.lemmatize(text) for text in texts]
+
+    texts.extend(lower_texts)
+    texts.extend(stem_texts)
+    texts.extend(lemm_texts)
+    return texts
+
+
+def extract_keywords(query_text, model):
+    prompt = f"###Instruction:Extract the important keywords which describe the context accurately.\n\nInput:{query_text}\n\n###Response:"
+    response = model.predict(prompt)
+    keywords = response.split(", ")
+    keywords = all_keywords_combs(keywords)
+    return keywords
+
 
 # Entity Extraction
 

utils/models.py

@@ -10,6 +10,7 @@ import streamlit_scrollable_textbox as stx
 import torch
 from InstructorEmbedding import INSTRUCTOR
 from sentence_transformers import SentenceTransformer
+from gradio_client import Client
 from tqdm import tqdm
 from transformers import (
     AutoModelForMaskedLM,
@@ -103,6 +104,12 @@ def get_instructor_embedding_model():
     return model
 
 
+@st.experimental_singleton
+def get_alpaca_model():
+    client = Client("https://awinml-alpaca-cpp.hf.space")
+    return client
+
+
 @st.experimental_memo
 def save_key(api_key):
     return api_key

utils/retriever.py

@@ -7,6 +7,7 @@ def query_pinecone_sparse(
     quarter,
     ticker,
     participant_type,
+    keywords=None,
     threshold=0.25,
 ):
     if participant_type == "Company Speaker":
@@ -33,6 +34,7 @@ def query_pinecone_sparse(
                     "Quarter": {"$in": ["Q1", "Q2", "Q3", "Q4"]},
                     "Ticker": {"$eq": ticker},
                     "QA_Flag": {"$eq": participant},
+                    "Keywords": {"$in": keywords}
                 },
                 include_metadata=True,
             )
@@ -54,6 +56,7 @@ def query_pinecone_sparse(
                     "Quarter": {"$eq": quarter},
                     "Ticker": {"$eq": ticker},
                     "QA_Flag": {"$eq": participant},
+                    "Keywords": {"$in": keywords}
                 },
                 include_metadata=True,
             )
@@ -68,6 +71,7 @@ def query_pinecone_sparse(
                 "Quarter": {"$eq": quarter},
                 "Ticker": {"$eq": ticker},
                 "QA_Flag": {"$eq": participant},
+                "Keywords": {"$in": keywords}
             },
             include_metadata=True,
         )
@@ -88,6 +92,7 @@ def query_pinecone(
     quarter,
     ticker,
     participant_type,
+    keywords=None,
     threshold=0.25,
 ):
     if participant_type == "Company Speaker":
@@ -113,6 +118,7 @@ def query_pinecone(
                     "Quarter": {"$in": ["Q1", "Q2", "Q3", "Q4"]},
                     "Ticker": {"$eq": ticker},
                     "QA_Flag": {"$eq": participant},
+                    "Keywords": {"$in": keywords}
                 },
                 include_metadata=True,
             )
@@ -133,6 +139,7 @@ def query_pinecone(
                     "Quarter": {"$eq": quarter},
                     "Ticker": {"$eq": ticker},
                     "QA_Flag": {"$eq": participant},
+                    "Keywords": {"$in": keywords}
                 },
                 include_metadata=True,
             )
@@ -146,6 +153,7 @@ def query_pinecone(
                 "Quarter": {"$eq": quarter},
                 "Ticker": {"$eq": ticker},
                 "QA_Flag": {"$eq": participant},
+                "Keywords": {"$in": keywords}
             },
             include_metadata=True,
         )