Upload 2 files

app.py

@@ -10,7 +10,8 @@ from utils import (
     clean_entities,
     create_dense_embeddings,
     create_sparse_embeddings,
-    extract_entities,
+    extract_quarter_year,
+    extract_ticker_spacy,
     format_query,
     get_flan_alpaca_xl_model,
     generate_alpaca_ner_prompt,
@@ -70,9 +71,12 @@ with col1:
 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)
+    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)
+    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

utils.py

@@ -5,6 +5,7 @@ import requests
 import openai
 import pandas as pd
 import spacy
+import spacy_transformers
 import streamlit_scrollable_textbox as stx
 import torch
 from sentence_transformers import SentenceTransformer
@@ -33,13 +34,17 @@ def get_data():
 
 @st.experimental_singleton
 def get_spacy_model():
-    return spacy.load("en_core_web_sm")
+    return spacy.load("en_core_web_trf")
 
 
 @st.experimental_singleton
 def get_flan_alpaca_xl_model():
-    model = AutoModelForSeq2SeqLM.from_pretrained("/home/user/app/models/flan-alpaca-xl/")
-    tokenizer = AutoTokenizer.from_pretrained("/home/user/app/models/flan-alpaca-xl/")
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        "/home/user/app/models/flan-alpaca-xl/"
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        "/home/user/app/models/flan-alpaca-xl/"
+    )
     return model, tokenizer
 
 
@@ -478,6 +483,7 @@ Answer:?"""
 
 # Entity Extraction
 
+
 def generate_alpaca_ner_prompt(query):
     prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Use the following guidelines to extract the entities representing the Company, Quarter, and Year in the sentence.
 
@@ -515,19 +521,27 @@ Company - Cisco, Quarter - none, Year - 2016
 ### Response:"""
     return prompt
 
+
 def generate_entities_flan_alpaca_inference_api(prompt):
     API_URL = "https://api-inference.huggingface.co/models/declare-lab/flan-alpaca-xl"
     API_TOKEN = st.secrets["hg_key"]
     headers = {"Authorization": f"Bearer {API_TOKEN}"}
     payload = {
         "inputs": prompt,
-        "parameters": {"do_sample": True, "temperature":0.1, "max_length":80},
-        "options": {"use_cache": False, "wait_for_model": True}
+        "parameters": {
+            "do_sample": True,
+            "temperature": 0.1,
+            "max_length": 80,
+        },
+        "options": {"use_cache": False, "wait_for_model": True},
     }
     try:
         data = json.dumps(payload)
+        # Key not used as headers=headers not passed
         response = requests.request("POST", API_URL, data=data)
-        output = json.loads(response.content.decode("utf-8"))[0]["generated_text"]
+        output = json.loads(response.content.decode("utf-8"))[0][
+            "generated_text"
+        ]
     except:
         output = ""
     print(output)
@@ -536,7 +550,7 @@ def generate_entities_flan_alpaca_inference_api(prompt):
 
 def generate_entities_flan_alpaca_checkpoint(model, tokenizer, prompt):
     model_inputs = tokenizer(prompt, return_tensors="pt")
-    input_ids = inputs["input_ids"]
+    input_ids = model_inputs["input_ids"]
     generation_output = model.generate(
         input_ids=input_ids,
         temperature=0.1,
@@ -547,9 +561,9 @@ def generate_entities_flan_alpaca_checkpoint(model, tokenizer, prompt):
     return output
 
 
-def format_entities_flan_alpaca(model_output):
+def format_entities_flan_alpaca(values):
     """
-    Extracts the text for each entity from the output generated by the 
+    Extracts the text for each entity from the output generated by the
     Flan-Alpaca model.
     """
     try:
@@ -560,22 +574,22 @@ def format_entities_flan_alpaca(model_output):
         year = None
     try:
         company = company_string.split(" - ")[1].lower()
-        company = None if company.lower() == 'none' else company
+        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
+        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
+        year = None if year.lower() == "none" else year
 
     except:
         year = None
-    
+
     print((company, quarter, year))
     return company, quarter, year
 
@@ -586,34 +600,27 @@ def extract_quarter_year(string):
     if year_match:
         year = year_match.group()
     else:
-        return None, None
+        year = None
 
     # Extract quarter from string
     quarter_match = re.search(r"Q\d", string)
     if quarter_match:
         quarter = "Q" + quarter_match.group()[1]
     else:
-        return None, None
+        quarter = None
 
     return quarter, year
 
 
-def extract_entities(query, model):
+def extract_ticker_spacy(query, model):
     doc = model(query)
     entities = {ent.label_: ent.text for ent in doc.ents}
+    print(entities.keys())
     if "ORG" in entities.keys():
         company = entities["ORG"].lower()
-        if "DATE" in entities.keys():
-            quarter, year = extract_quarter_year(entities["DATE"])
-            return company, quarter, year
-        else:
-            return company, None, None
     else:
-        if "DATE" in entities.keys():
-            quarter, year = extract_quarter_year(entities["DATE"])
-            return None, quarter, year
-        else:
-            return None, None, None
+        company = None
+    return company
 
 
 def clean_entities(company, quarter, year):