Update app.py

app.py

@@ -1,17 +1,21 @@
 """Demo for NER4OPT."""
 import os
 import warnings
-
-warnings.filterwarnings("ignore", category=DeprecationWarning)
-
 import streamlit as st
 import ner4opt
 from ner4opt.utils import preprocess, spacy_tokenize_sentence
-
 from spacy import displacy
 
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+
 HTML_WRAPPER = """<div style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1rem; margin-bottom: 2.5rem">{}</div>"""
 
+# Define your example prompts here for easy management
+EXAMPLE_PROMPTS = {
+    "Example 1️⃣: Asset Allocation": "Cautious Asset Investment has a total of $150,000 to manage and decides to invest it in money market fund, which yields a 2% return as well as in foreign bonds, which gives and average rate of return of 10.2%. Internal policies require PAI to diversify the asset allocation so that the minimum investment in money market fund is 40% of the total investment. Due to the risk of default of foreign countries, no more than 40% of the total investment should be allocated to foreign bonds. How much should the Cautious Asset Investment allocate in each asset so as to maximize its average return?",
+    "Example 2️⃣: Production Optimization": "A factory produces two types of goods, A and B. Producing one unit of good A requires 2 hours of labor and 1 unit of raw material. Producing one unit of good B requires 3 hours of labor and 2 units of raw material. The factory has a total of 100 labor hours and 50 units of raw material available. If the profit from selling one unit of good A is $10 and from good B is $15, how many units of each good should be produced to maximize the total profit?",
+    "Example 3️⃣: Transportation Problem": "A company has two warehouses and three retail stores. Warehouse 1 has 500 units of a product and Warehouse 2 has 700 units. Store 1 requires 300 units, Store 2 requires 400 units, and Store 3 requires 500 units. The cost of shipping one unit from Warehouse 1 to Store 1 is $2, to Store 2 is $3, and to Store 3 is $4. The cost of shipping from Warehouse 2 to Store 1 is $5, to Store 2 is $2, and to Store 3 is $3. How should the company ship the products to minimize the total shipping cost?"
+}
 
 @st.cache_resource
 def load_models():
@@ -22,27 +26,41 @@ def load_models():
     hybrid_ner_model = ner4opt.Ner4Opt("hybrid")
     return lexical_ner_model, lexical_plus_ner_model, semantic_ner_model, hybrid_ner_model
 
+def set_example_text(example_key):
+    """Callback to set the text area value based on the selected example."""
+    st.session_state.text_input = EXAMPLE_PROMPTS[example_key]
 
 def main():
-
     st.title("""Ner4Opt: Named Entity Recognition for Optimization""")
-
     st.markdown("""Given an optimization problem in natural language, Ner4Opt extracts optimization related entities from free-form text. The source code for Ner4Opt is available at https://github.com/skadio/ner4opt""")
 
+    # Initialize session state for the text input if it doesn't exist
+    if "text_input" not in st.session_state:
+        st.session_state.text_input = EXAMPLE_PROMPTS["Example 1️⃣: Asset Allocation"]
+
+    # --- Add the example prompts section ---
+    st.subheader("Try with an example:")
+    cols = st.columns(len(EXAMPLE_PROMPTS))
+    for i, (key, value) in enumerate(EXAMPLE_PROMPTS.items()):
+        with cols[i]:
+            if st.button(key):
+                set_example_text(key)
+    
+    # ----------------------------------------
+    
     option = st.sidebar.selectbox(
         'Select a lexical, semantic, or hybrid model for extracting entities',
         ('Lexical', 'Lexical Plus', 'Semantic', 'Hybrid'), index=3)
 
     text = st.text_area(
         "Text",
-        "Cautious Asset Investment has a total of $150,000 to manage and decides to invest it in money market fund, which yields a 2% return as well as in foreign bonds, which gives and average rate of return of 10.2%. Internal policies require PAI to diversify the asset allocation so that the minimum investment in money market fund is 40% of the total investment. Due to the risk of default of foreign countries, no more than 40% of the total investment should be allocated to foreign bonds. How much should the Cautious Asset Investment allocate in each asset so as to maximize its average return?"
+        value=st.session_state.text_input, # Use the session state value here
     )
     text = text.strip()
     if text == "":
         st.write("Please write a valid sentence.")
 
-    lexical_ner_model, lexical_plus_ner_model, semantic_ner_model, hybrid_ner_model = load_models(
-    )
+    lexical_ner_model, lexical_plus_ner_model, semantic_ner_model, hybrid_ner_model = load_models()
 
     # get entities
     if option == "Lexical":
@@ -74,6 +92,5 @@ def main():
     html_ner = html_ner.replace("\n", " ")
     st.write(HTML_WRAPPER.format(html_ner), unsafe_allow_html=True)
 
-
 if __name__ == '__main__':
-    main()
+    main()