Upload 2 files

app.py

@@ -13,6 +13,7 @@ from utils import (
     extract_entities,
     format_query,
     get_flan_alpaca_xl_model,
+    generate_alpaca_ner_prompt,
     generate_entities_flan_alpaca,
     format_entities_flan_alpaca,
     generate_flant5_prompt_instruct_chunk_context,
@@ -56,7 +57,7 @@ with st.sidebar:
     ner_choice = st.selectbox("Select NER Model", ["Alpaca", "Spacy"])
 
 if ner_choice == "Alpaca":
-    ner_model = get_flan_alpaca_xl_model()
+    ner_model, ner_tokenizer = get_flan_alpaca_xl_model()
 else:
     ner_model = get_spacy_model()
 
@@ -68,7 +69,8 @@ with col1:
     )
 
 if ner_choice == "Alpaca":
-    entity_text = generate_entities_flan_alpaca(ner_model)
+    ner_prompt = generate_alpaca_ner_prompt(query_text)
+    entity_text = generate_entities_flan_alpaca(ner_model, ner_tokenizer, ner_prompt)
     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)

utils.py

@@ -36,7 +36,9 @@ def get_spacy_model():
 
 @st.experimental_singleton
 def get_flan_alpaca_xl_model():
-    return pipeline(model="./models/flan-alpaca-xl")
+    model = AutoModelForSeq2SeqLM("./models/flan-alpaca-xl")
+    tokenizer = AutoTokenizer("./models/flan-alpaca-xl")
+    return model, tokenizer
 
 
 # Initialize models from HuggingFace
@@ -474,10 +476,55 @@ 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 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.
+
+### Instruction:
+- The output should be in the form "Company - Value, Quarter - Value, Year - Value".
+- The output should be in the form "Company - None, Quarter - None, Year - None", if no entities are found.
+- Only use entities that exist in the final sentence.
+- If Company cannot be found in the sentence, return "none" for that entity.
+- If Quarter cannot be found in the sentence, return "none" for that entity.
+- If Year cannot be found in the sentence, return "none" for that entity.
+- If there is ambiguity finding the entity, return "none" for that entity.
+
+### Input:
+
+What was discussed regarding Wearables revenue performance in Apple's Q3 2023 earnings call?
+Company - Apple, Quarter - Q3, Year - 2023
+
+How has the growth in Q1 been for the PC market as seen by AMD?
+Company - AMD, Quarter - Q1, Year - none
+
+What was the long term view on GOOGL's cloud business growth as discussed in their earnings call?
+Company - Google, Quarter - none, Year - none
+
+What is Nvidia's visibility in the data center business in 2020?
+Company - Nvidia, Quarter - none, Year - 2020
+
+What are the opportunities and challenges in the Indian market that Amazon is facing?
+Company - Amazon, Quarter - none, Year - none
+
+What did the Analysts ask about CSCO's cybersecurity business in the earnings call?
+Company - Cisco, Quarter - none, Year - none
+
+
+{query}
+### Response:"""
+    return prompt
+
+
+def generate_entities_flan_alpaca(model, tokenizer, prompt):
+    model_inputs = tokenizer(prompt, return_tensors="pt")
+    input_ids = inputs["input_ids"]
+    generation_output = model.generate(
+        input_ids=input_ids,
+        temperature=0.1,
+        top_p=0.5,
+        max_new_tokens=1024,
+    )
+    output = tokenizer.decode(generation_output[0], skip_special_tokens=True)
+    return output
 
 
 def format_entities_flan_alpaca(model_output):