app.py

"""Demo for NER4OPT."""
import os
import warnings
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>"""

# Configuration: Change this number to control how many examples to show by default
DEFAULT_EXAMPLES_TO_SHOW = 3
EXAMPLES_PER_ROW = 3

# 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?",
    
    "Example 4️⃣: Asset Investment": "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 5️⃣: Production Planning": "A manufacturing company produces two types of products: Product A and Product B. Each Product A requires 3 hours of labor and 2 units of raw material, while each Product B requires 2 hours of labor and 4 units of raw material. The company has 240 hours of labor and 200 units of raw material available per week. Product A generates a profit of $50 per unit and Product B generates a profit of $40 per unit. How many units of each product should the company produce to maximize weekly profit?",
    
    "Example 6️⃣: Transportation Problem": "A company has three warehouses (W1, W2, W3) with supplies of 100, 150, and 200 units respectively. They need to ship goods to four retail stores (S1, S2, S3, S4) with demands of 80, 120, 90, and 160 units respectively. The transportation costs per unit from each warehouse to each store are given in a cost matrix. Find the optimal shipping plan that minimizes the total transportation cost while satisfying all supply and demand constraints.",
    
    "Example 7️⃣: Diet Optimization": "A nutritionist is designing a diet plan that includes chicken, fish, and vegetables. Each serving of chicken provides 25g protein, 200 calories, and costs $3. Each serving of fish provides 30g protein, 150 calories, and costs $4. Each serving of vegetables provides 5g protein, 50 calories, and costs $1. The diet must provide at least 100g protein and at most 800 calories per day. What combination of servings minimizes the daily cost while meeting nutritional requirements?",
    
    "Example 8️⃣: Facility Location": "A retail chain wants to open new distribution centers to serve 5 cities. They have identified 3 potential locations for distribution centers. Each location has a fixed opening cost and variable shipping costs to each city. The company wants to minimize the total cost (fixed costs plus shipping costs) while ensuring each city is served by exactly one distribution center. Which distribution centers should be opened and how should cities be assigned to minimize total cost?",
    
    "Example 9️⃣: Resource Allocation": "A project manager has 4 teams available to work on 3 different projects. Each team has different productivity rates for each project and different hourly costs. Project A requires 100 hours, Project B requires 150 hours, and Project C requires 200 hours. Each team can work a maximum of 160 hours this month. How should the manager assign teams to projects to minimize total cost while completing all projects on time?",

    "Example 🔟: Diet Optimization": "A diet plan requires at least 3000 calories, 50 grams of protein, and 70 grams of fat. Food X provides 600 calories, 20g protein, and 10g fat per serving. Food Y provides 500 calories, 10g protein, and 20g fat per serving. How many servings of each food should be consumed to meet the requirements at minimum cost?",
}

@st.cache_resource
def load_models():
    """Load and cache NER4OPT models."""
    lexical_ner_model = ner4opt.Ner4Opt("lexical")
    lexical_plus_ner_model = ner4opt.Ner4Opt("lexical_plus")
    semantic_ner_model = ner4opt.Ner4Opt("semantic")
    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 display_example_buttons(show_all=False):
    """Display example buttons in organized rows."""
    example_items = list(EXAMPLE_PROMPTS.items())
    
    if not show_all:
        example_items = example_items[:DEFAULT_EXAMPLES_TO_SHOW]
    
    # Display examples in rows
    for i in range(0, len(example_items), EXAMPLES_PER_ROW):
        cols = st.columns(EXAMPLES_PER_ROW)
        
        for j, (key, value) in enumerate(example_items[i:i+EXAMPLES_PER_ROW]):
            with cols[j]:
                if st.button(key, key=f"example_btn_{i+j}", use_container_width=True):
                    set_example_text(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"]
    
    # Initialize session state for showing all examples
    if "show_all_examples" not in st.session_state:
        st.session_state.show_all_examples = False

    # --- Add the example prompts section ---
    st.subheader("📝 Try with an example:")
    
    # Toggle button for showing more/fewer examples
    total_examples = len(EXAMPLE_PROMPTS)
    if total_examples > DEFAULT_EXAMPLES_TO_SHOW:
        col1, col2 = st.columns([3, 1])
        with col2:
            if st.session_state.show_all_examples:
                if st.button(f"Show fewer ({DEFAULT_EXAMPLES_TO_SHOW})", key="toggle_examples"):
                    st.session_state.show_all_examples = False
                    st.rerun()
            else:
                if st.button(f"Show all ({total_examples})", key="toggle_examples"):
                    st.session_state.show_all_examples = True
                    st.rerun()
    
    # Display the example buttons
    display_example_buttons(st.session_state.show_all_examples)
    
    st.markdown("---")
    # ----------------------------------------
    
    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 (Tip: Use Ctrl+Enter to process)",
        value=st.session_state.text_input, # Use the session state value here
        height=120,
        help="Enter your optimization problem description or select an example above."
    )
    
    text = text.strip()
    if text == "":
        st.warning("Please write a valid sentence or select an example above.")
        return

    try:
        with st.spinner("Processing text..."):
            lexical_ner_model, lexical_plus_ner_model, semantic_ner_model, hybrid_ner_model = load_models()

            # get entities
            if option == "Lexical":
                predicted_entities = lexical_ner_model.get_entities(text)
            elif option == "Lexical Plus":
                predicted_entities = lexical_plus_ner_model.get_entities(text)
            elif option == "Semantic":
                predicted_entities = semantic_ner_model.get_entities(text)
            elif option == "Hybrid":
                predicted_entities = hybrid_ner_model.get_entities(text)

            # re-format to match SpaCy, as it uses label instead of entity_group
            # entity_group follows huggingface design and is compatible with most NER interfaces
            entities_formatted = []
            for tag in predicted_entities:
                entities_formatted.append({
                    "start": tag['start'],
                    "end": tag['end'],
                    "label": tag['entity_group'],
                    "score": tag['score']
                })

            ner_for_display = [{
                "text": spacy_tokenize_sentence(preprocess(text)),
                "ents": entities_formatted,
                "title": "Named Entities"
            }]
            html_ner = displacy.render(ner_for_display, style="ent", manual=True)
            html_ner = html_ner.replace("\n", " ")
            st.write(HTML_WRAPPER.format(html_ner), unsafe_allow_html=True)
            
            # Show entity count if any entities found
            if entities_formatted:
                st.success(f"Found {len(entities_formatted)} entities using {option} model")
            else:
                st.info("No entities detected in the text.")
                
    except Exception as e:
        st.error(f"An error occurred: {str(e)}")

if __name__ == '__main__':
    main()