Upload 2 files

app.py

@@ -12,15 +12,18 @@ from utils import (
     create_sparse_embeddings,
     extract_entities,
     format_query,
+    get_flan_alpaca_xl_model,
+    generate_entities_flan_alpaca,
+    format_entities_flan_alpaca,
     generate_flant5_prompt_instruct_chunk_context,
-    generate_flant5_prompt_instruct_complete_context,
     generate_flant5_prompt_instruct_chunk_context_single,
-    generate_flant5_prompt_summ_chunk_context_single,
+    generate_flant5_prompt_instruct_complete_context,
     generate_flant5_prompt_summ_chunk_context,
-    generate_text_flan_t5,
-    generate_gpt_prompt,
+    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_text_flan_t5,
     get_context_list_prompt,
     get_data,
     get_flan_t5_model,
@@ -49,7 +52,13 @@ st.write(
 col1, col2 = st.columns([3, 3], gap="medium")
 
 
-spacy_model = get_spacy_model()
+with st.sidebar:
+    ner_choice = st.selectbox("Select NER Model", ["Alpaca", "Spacy"])
+
+if ner_choice == "Alpaca":
+    ner_model = get_flan_alpaca_xl_model()
+else:
+    ner_model = get_spacy_model()
 
 with col1:
     st.subheader("Question")
@@ -58,7 +67,12 @@ with col1:
         value="What was discussed regarding Wearables revenue performance?",
     )
 
-company_ent, quarter_ent, year_ent = extract_entities(query_text, spacy_model)
+if ner_choice == "Alpaca":
+    entity_text = generate_entities_flan_alpaca(ner_model)
+    company_ent, quarter_ent, year_ent = format_entities_flan_alpaca(entity_text)
+else:
+    company_ent, quarter_ent, year_ent = extract_entities(query_text, ner_model)
+
 ticker_index, quarter_index, year_index = clean_entities(
     company_ent, quarter_ent, year_ent
 )
@@ -251,7 +265,9 @@ if decoder_model == "GPT3 - (text-davinci-003)":
 
 
 elif decoder_model == "T5":
-    prompt = generate_flant5_prompt_instruct_complete_context(query_text, context_list)
+    prompt = generate_flant5_prompt_instruct_complete_context(
+        query_text, context_list
+    )
     t5_pipeline = get_t5_model()
     output_text = []
     with col2:
@@ -275,7 +291,8 @@ elif decoder_model == "FLAN-T5":
     output_text = []
     with col2:
         prompt_type = st.selectbox(
-            "Select prompt type", ["Complete Text QA", "Chunkwise QA", "Chunkwise Summarize"]
+            "Select prompt type",
+            ["Complete Text QA", "Chunkwise QA", "Chunkwise Summarize"],
         )
         if prompt_type == "Complete Text QA":
             prompt = generate_flant5_prompt_instruct_complete_context(
@@ -300,23 +317,37 @@ elif decoder_model == "FLAN-T5":
             submitted = st.form_submit_button("Submit")
             if submitted:
                 if prompt_type == "Complete Text QA":
-                    output_text_string = generate_text_flan_t5(flan_t5_model, flan_t5_tokenizer, prompt)
+                    output_text_string = generate_text_flan_t5(
+                        flan_t5_model, flan_t5_tokenizer, prompt
+                    )
                     st.subheader("Answer:")
                     st.write(output_text_string)
                 elif prompt_type == "Chunkwise QA":
                     for context_text in context_list:
-                        model_input = generate_flant5_prompt_instruct_chunk_context_single(query_text, context_text)
+                        model_input = generate_flant5_prompt_instruct_chunk_context_single(
+                            query_text, context_text
+                        )
                         output_text.append(
-                            generate_text_flan_t5(flan_t5_model, flan_t5_tokenizer, model_input))
+                            generate_text_flan_t5(
+                                flan_t5_model, flan_t5_tokenizer, model_input
+                            )
+                        )
                     st.subheader("Answer:")
                     for text in output_text:
                         if "(iii)" not in text:
                             st.markdown(f"- {text}")
                 elif prompt_type == "Chunkwise Summarize":
                     for context_text in context_list:
-                        model_input = generate_flant5_prompt_summ_chunk_context_single(query_text, context_text)
+                        model_input = (
+                            generate_flant5_prompt_summ_chunk_context_single(
+                                query_text, context_text
+                            )
+                        )
                         output_text.append(
-                            generate_text_flan_t5(flan_t5_model, flan_t5_tokenizer, model_input))
+                            generate_text_flan_t5(
+                                flan_t5_model, flan_t5_tokenizer, model_input
+                            )
+                        )
                     st.subheader("Answer:")
                     for text in output_text:
                         if "(iii)" not in text:

utils.py

@@ -2,7 +2,6 @@ import re
 
 import openai
 import pandas as pd
-import pinecone
 import spacy
 import streamlit_scrollable_textbox as stx
 import torch
@@ -12,11 +11,12 @@ from transformers import (
     AutoModelForMaskedLM,
     AutoModelForSeq2SeqLM,
     AutoTokenizer,
+    T5ForConditionalGeneration,
+    T5Tokenizer,
     pipeline,
 )
-from transformers import T5Tokenizer, T5ForConditionalGeneration
-
 
+import pinecone
 import streamlit as st
 
 
@@ -34,6 +34,11 @@ def get_spacy_model():
     return spacy.load("en_core_web_sm")
 
 
+@st.experimental_singleton
+def get_flan_alpaca_xl_model():
+    return pipeline(model="declare-lab/flan-alpaca-xl")
+    
+
 # Initialize models from HuggingFace
 
 
@@ -469,6 +474,42 @@ Answer:?"""
 
 # Entity Extraction
 
+def generate_entities_flan_alpaca(model):
+    output = model(prompt, max_length=512, temperature=0.1)
+    generated_text = output[0]["generated_text"]
+    return generated_text
+
+
+def format_entities_flan_alpaca(model_output):
+    """
+    Extracts the text for each entity from the output generated by the 
+    Flan-Alpaca model.
+    """
+    try:
+        company_string, quarter_string, year_string = values.split(", ")
+    except:
+        company = None
+        quarter = None
+        year = None
+    try:
+        company = company_string.split(" - ")[1].lower()
+        company = None if company.lower() == 'none' else company
+    except:
+        company = None
+    try:
+        quarter = quarter_string.split(" - ")[1]
+        quarter = None if quarter.lower() == 'none' else quarter
+
+    except:
+        quarter = None
+    try:
+        year = year_string.split(" - ")[1]
+        year = None if year.lower() == 'none' else year
+
+    except:
+        year = None
+    return company, quarter, year
+
 
 def extract_quarter_year(string):
     # Extract year from string