Upload 2 files

app.py

@@ -14,7 +14,8 @@ from utils import (
     format_query,
     get_flan_alpaca_xl_model,
     generate_alpaca_ner_prompt,
-    generate_entities_flan_alpaca,
+    generate_entities_flan_alpaca_checkpoint,
+    generate_entities_flan_alpaca_inference_api,
     format_entities_flan_alpaca,
     generate_flant5_prompt_instruct_chunk_context,
     generate_flant5_prompt_instruct_chunk_context_single,
@@ -56,9 +57,7 @@ col1, col2 = st.columns([3, 3], gap="medium")
 with st.sidebar:
     ner_choice = st.selectbox("Select NER Model", ["Alpaca", "Spacy"])
 
-if ner_choice == "Alpaca":
-    ner_model, ner_tokenizer = get_flan_alpaca_xl_model()
-else:
+if ner_choice == "Spacy":
     ner_model = get_spacy_model()
 
 with col1:
@@ -70,7 +69,7 @@ with col1:
 
 if ner_choice == "Alpaca":
     ner_prompt = generate_alpaca_ner_prompt(query_text)
-    entity_text = generate_entities_flan_alpaca(ner_model, ner_tokenizer, ner_prompt)
+    entity_text = generate_entities_flan_alpaca_inference_api(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

@@ -1,4 +1,6 @@
 import re
+import json
+import requests
 
 import openai
 import pandas as pd
@@ -513,8 +515,20 @@ Company - Cisco, Quarter - none, Year - none
 ### Response:"""
     return prompt
 
+def generate_entities_flan_alpaca_inference_api(prompt):
+    API_URL = "https://api-inference.huggingface.co/models/declare-lab/flan-alpaca-xl"
+    payload = {
+        "inputs": prompt,
+        "parameters": {"do_sample": True, "temperature":0.1, "max_length":80},
+        "options": {"use_cache": True, "wait_for_model": True}
+    }
+    data = json.dumps(payload)
+    response = requests.request("POST", API_URL, data=data)
+    output = json.loads(response.content.decode("utf-8"))[0]["generated_text"]
+    return output
+
 
-def generate_entities_flan_alpaca(model, tokenizer, prompt):
+def generate_entities_flan_alpaca_checkpoint(model, tokenizer, prompt):
     model_inputs = tokenizer(prompt, return_tensors="pt")
     input_ids = inputs["input_ids"]
     generation_output = model.generate(