Update app.py

app.py

@@ -14,8 +14,8 @@ from transformers import AutoTokenizer, AutoModelForTokenClassification
 def load_etf_data():
     """
     Load ETF data with a persistent 'doc' column.
-    - On first run: reads raw CSV, computes 'doc', saves enriched CSV.
-    - On subsequent runs: loads enriched CSV directly.
+    - First run: reads raw CSV, computes 'doc', saves enriched CSV.
+    - Subsequent runs: loads enriched CSV directly.
     """
     enriched_path = "etf_general_info_enriched_doc_added.csv"
     raw_path      = "etf_general_info_enriched.csv"
@@ -27,29 +27,21 @@ def load_etf_data():
         df_info["doc"] = df_info.apply(modules.make_doc_text, axis=1)
         df_info.to_csv(enriched_path, index=False)
 
-    # Split into DataFrame and ticker list
     df_etf, available_tickers = modules.set_etf_data(df_info)
-
-    # Load other supporting DataFrames
     df_analyst_report       = pd.read_csv("etf_analyst_report_full.csv")
-    df_annual_return_master = (
-        pd.read_csv("annual_return.csv").rename(columns={"ticker": "Ticker"})
-    )
+    df_annual_return_master = pd.read_csv("annual_return.csv").rename(columns={"ticker": "Ticker"})
     return df_etf, df_analyst_report, available_tickers, df_annual_return_master
 
 @st.cache_resource(show_spinner=False)
 def build_search_resources():
     """
-    Loads (or builds) SentenceTransformer + FAISS index + ticker list.
-    - On first run: computes embeddings, builds index, writes to disk.
-    - On subsequent runs: loads FAISS index from disk.
+    Load or build SentenceTransformer + FAISS index + ticker list.
+    - First run: computes embeddings, builds index, writes to disk.
+    - Subsequent runs: loads FAISS index from disk.
     """
     df_etf, *_ = load_etf_data()
-
-    # Load SentenceTransformer
-    repo_name = "hskwon7/paraphrase-MiniLM-L6-v2-ft-for-etf-semantic-search"
-    model     = SentenceTransformer(repo_name)
-
+    repo_name  = "hskwon7/paraphrase-MiniLM-L6-v2-ft-for-etf-semantic-search"
+    model      = SentenceTransformer(repo_name)
     ticker_list = df_etf["Ticker"].tolist()
 
     index_path = "etf_faiss.index"
@@ -68,28 +60,25 @@ def build_search_resources():
 @st.cache_resource(show_spinner=False)
 def load_ner_models():
     """
-    Loads two NER models for ticker extraction, and builds the valid ticker set.
+    Loads two NER models for ticker extraction and builds valid_ticker_set.
     """
     repo1 = "hskwon7/distilbert-base-uncased-for-etf-ticker"
     repo2 = "hskwon7/albert-base-v2-for-etf-ticker"
-
-    tok1 = AutoTokenizer.from_pretrained(repo1)
-    m1   = AutoModelForTokenClassification.from_pretrained(repo1)
-    tok2 = AutoTokenizer.from_pretrained(repo2)
-    m2   = AutoModelForTokenClassification.from_pretrained(repo2)
+    tok1, m1 = AutoTokenizer.from_pretrained(repo1), AutoModelForTokenClassification.from_pretrained(repo1)
+    tok2, m2 = AutoTokenizer.from_pretrained(repo2), AutoModelForTokenClassification.from_pretrained(repo2)
 
     df_etf, *_ = load_etf_data()
     valid_ticker_set = set(t.upper() for t in df_etf["Ticker"].unique())
 
     return (tok1, m1), (tok2, m2), valid_ticker_set
 
-# ─── INITIALIZE CACHED RESOURCES ─────────────────────────────────────────────
+# ─── INITIALIZE ─────────────────────────────────────────────────────────────
 
 df_etf, df_analyst_report, available_tickers, df_annual_return_master = load_etf_data()
 s2_model, faiss_index, etf_list = build_search_resources()
 (tok1, m1), (tok2, m2), valid_ticker_set = load_ner_models()
 
-# ─── CORE SEARCH & EXTRACTION ─────────────────────────────────────────────────
+# ─── CORE SEARCH & EXTRACTION ───────────────────────────────────────────────
 
 def semantic_search(query: str, top_k: int = 100):
     q_emb = s2_model.encode([query], convert_to_numpy=True)
@@ -100,54 +89,94 @@ def semantic_search(query: str, top_k: int = 100):
 def ensemble_ticker_extraction(query: str):
     preds = set()
     for tok, mdl in ((tok1, m1), (tok2, m2)):
-        enc    = tok(query, return_tensors="pt")
+        enc = tok(query, return_tensors="pt")
         with torch.no_grad():
             logits = mdl(**enc).logits
-        ids    = logits.argmax(dim=-1)[0].tolist()
-        toks   = tok.convert_ids_to_tokens(enc["input_ids"][0])
-        labs   = [mdl.config.id2label[i] for i in ids]
+        ids  = logits.argmax(dim=-1)[0].tolist()
+        toks = tok.convert_ids_to_tokens(enc["input_ids"][0])
+        labs = [mdl.config.id2label[i] for i in ids]
         preds |= modules.extract_valid_tickers(toks, labs, tok, valid_ticker_set)
     return preds
 
-# ─── HELPERS ────────────────────────���───────────────────────────────────────
+# ─── UI HELPERS ──────────────────────────────────────────────────────────────
 
-def display_sample_query_box(task: str):
+def display_sample_query_boxes(key_prefix=""):
     sample_queries = {
         "search_etf": {
-            "title": "ETF Search",
-            "description": "Explore ETFs based on dividend, expense ratio, sector, etc.",
-            "examples": ['Find me technology ETFs', 'Show me Cryptocurrency ETFs']
+            "title": "AI ETF Search",
+            "description": "Explore ETFs by dividend, expense ratio, sector, etc.",
+            "query": ['Find me some high-dividend ETFs in the tech sector', 'Any ETFs holding Apple?']
         },
         "comparison": {
-            "title": "ETF Comparison",
-            "description": "Compare two ETFs side by side on performance, risk, etc.",
-            "examples": ["QQQ vs. SPY", "Compare performance of QQQ with SPY"]
+            "title": "ETF Performance Comparison",
+            "description": "Compare two ETFs side by side on performance and risk.",
+            "query": "I'd like to compare performance of QQQ with SPY."
         },
         "portfolio_projection": {
-            "title": "ETF Portfolio",
-            "description": "Project a multi-ETF portfolio out over 30 years.",
-            "examples": ["SPY, GLD, BND", "I want to invest in SCHD and IAU"]
-        }
+            "title": "Portfolio Projection",
+            "description": "Project a portfolio with multiple ETFs over a set number of years.",
+            "query": "I want to invest in SPY, QQQ, SCHD, and IAU for 20 years."
+        },
     }
-    details = sample_queries[task]
-    box_html = f"""
-    <div style='border:1px solid #ddd;padding:1rem;border-radius:8px;'>
-      <h4>{details['title']}</h4>
-      <p style='color:#555;margin-bottom:.5rem;'>{details['description']}</p>
-      <p style='font-style:italic;color:#333;'>
-        Examples:<br/>{'<br/>'.join(details['examples'])}
-      </p>
-    </div>
-    """
-    st.markdown(box_html, unsafe_allow_html=True)
-    if st.button("Try this app", key=f"try_{task}"):
-        page_map = {
-            "search_etf":       "ETF Search",
-            "comparison":       "ETF Comparison",
-            "portfolio_projection": "ETF Portfolio"
-        }
-        st.session_state["page"] = page_map[task]
-        st.experimental_rerun()
+
+    cols = st.columns(len(sample_queries))
+    title_h = "60px"
+    desc_h  = "100px"
+    query_h = "80px"
+
+    st.markdown("""
+    <style>
+      .small-link-button {
+        font-size: 12px;
+        color: #0073e6;
+        text-decoration: underline;
+        cursor: pointer;
+      }
+    </style>
+    """, unsafe_allow_html=True)
+
+    for idx, (key, details) in enumerate(sample_queries.items()):
+        with cols[idx]:
+            box_html = f"""
+            <div style="
+              width: 100%; height: 300px; border: 1px solid #ddd;
+              border-radius: 10px; padding: 15px; margin: 10px auto;
+              background-color: #fff; box-shadow: 2px 2px 8px rgba(0,0,0,0.1);
+              display: flex; flex-direction: column; justify-content: space-between;
+            ">
+              <div style="height: {title_h}; text-align: center;">
+                <p style="margin:5px 0; color:#2c3e50; font-size:16px;">
+                  <b>{details['title']}</b>
+                </p>
+              </div>
+              <div style="height: {desc_h}; text-align: center; overflow:auto;">
+                <p style="margin:5px 0; color:#7f8c8d; font-size:14px;">
+                  {details['description']}
+                </p>
+              </div>
+              <div style="height: {query_h}; text-align: center; overflow:auto;">
+                <p style="margin:5px 0; font-style:italic; color:#34495e; font-size:13px;">
+                  {"<br>".join(f'&quot;{q}&quot;' for q in details['query']) 
+                     if isinstance(details['query'], list)
+                     else f'&quot;{details["query"]}&quot;'}
+                </p>
+              </div>
+            </div>
+            """
+            st.markdown(box_html, unsafe_allow_html=True)
+
+            # center the button
+            l, c, r = st.columns([1,2,1])
+            with c:
+                if st.button("Try this app", key=key_prefix + key):
+                    # just switch page
+                    page_map = {
+                        "search_etf":       "ETF Search",
+                        "comparison":       "ETF Comparison",
+                        "portfolio_projection": "ETF Portfolio"
+                    }
+                    st.session_state["page"] = page_map[key]
+                    st.rerun()
 
 def display_chat_history(task: str):
     hist = st.session_state.get(f"all_chat_history_{task}", [])
@@ -161,122 +190,110 @@ def display_chat_history(task: str):
 def process_query(task: str, query: str):
     top_k, top_n = 100, 30
 
-    if task == 'search_etf':
+    if task == "search_etf":
         with st.spinner("Searching ETFs..."):
-            fetched = semantic_search(query, top_k=top_k)
+            fetched = semantic_search(query, top_k)
             df_out  = modules.get_etf_recommendations_from_list(
-                fetched, df_etf, top_n=top_n
+                fetched, modules.get_cols_to_display(), df_etf, top_n
             )
+        msg = f"{len(df_out)} ETFs found."
         st.session_state[f"all_chat_history_{task}"].append(
-            modules.form_d_chat_history(
-                result_id=str(uuid.uuid4()),
-                response=f"{len(df_out)} ETFs found.",
-                task=task,
-                df=df_out
-            )
+            modules.form_d_chat_history(str(uuid.uuid4()), msg, task, df=df_out)
         )
 
-    elif task == 'comparison':
+    elif task == "comparison":
         with st.spinner("Running comparison..."):
-            tickers = ensemble_ticker_extraction(query)
-            if len(tickers) != 2:
-                response, fig, df_out = (
-                    "Please specify exactly two tickers.", None, None
-                )
+            tk = ensemble_ticker_extraction(query)
+            if len(tk) != 2:
+                resp, fig, df_out = "Please specify exactly two tickers.", None, None
             else:
                 df_out = modules.get_etf_recommendations_from_list(
-                    tickers,
-                    df_etf, top_n=2
+                    [(t, None) for t in tk],
+                    modules.get_cols_to_display(), df_etf, top_n=2
                 )
-                fig = modules.compare_etfs_interactive(*tickers)
-                response = f"Compared {tickers[0]} vs. {tickers[1]}."
+                fig  = modules.compare_etfs_interactive(*tk)
+                resp = f"Compared {tk[0]} vs. {tk[1]}."
         st.session_state[f"all_chat_history_{task}"].append(
-            modules.form_d_chat_history(
-                result_id=str(uuid.uuid4()),
-                response=response,
-                task=task,
-                fig=fig,
-                df=df_out
-            )
+            modules.form_d_chat_history(str(uuid.uuid4()), resp, task, fig=fig, df=df_out)
         )
 
-    elif task == 'portfolio_projection':
+    elif task == "portfolio_projection":
         with st.spinner("Projecting portfolio..."):
-            fetched = semantic_search(query, top_k=top_k)
-            df_port = modules.run_portfolio_analysis(
-                fetched, df_etf, df_annual_return_master
-            )
+            fetched = semantic_search(query, top_k)
+            df_port = modules.run_portfolio_analysis(fetched, df_etf, df_annual_return_master)
             fig     = modules.portfolio_interactive_chart(df_port)
-            response = "30-year projection generated."
+            resp    = "30-year projection generated."
         st.session_state[f"all_chat_history_{task}"].append(
-            modules.form_d_chat_history(
-                result_id=str(uuid.uuid4()),
-                response=response,
-                task=task,
-                fig=fig
-            )
+            modules.form_d_chat_history(str(uuid.uuid4()), resp, task, fig=fig)
         )
 
-def display_explore_etfs_chat():
-    st.markdown("""
-        <style>
-        .main .block-container {
-            max-width: 90% !important;
-            padding: 1rem 2rem;
+def main():
+    st.set_page_config(layout="wide")
+
+    # init state
+    if "page" not in st.session_state:
+        st.session_state["page"] = "Home"
+    if "user_query" not in st.session_state:
+        st.session_state["user_query"] = ""
+    if "auto_query_sent" not in st.session_state:
+        st.session_state["auto_query_sent"] = False
+    for t in ["search_etf","comparison","portfolio_projection"]:
+        st.session_state.setdefault(f"all_chat_history_{t}", [])
+
+    # sidebar
+    st.sidebar.title("ETF Assistant")
+    if st.sidebar.button("Home"):
+        st.session_state["page"] = "Home"
+    if st.sidebar.button("ETF Search"):
+        st.session_state["page"] = "ETF Search"
+    if st.sidebar.button("ETF Comparison"):
+        st.session_state["page"] = "ETF Comparison"
+    if st.sidebar.button("ETF Portfolio"):
+        st.session_state["page"] = "ETF Portfolio"
+
+    # render
+    page = st.session_state["page"]
+    if page == 'Home':
+        st.title("ETF Assistant")
+    else:
+        st.title(page)
+
+    # Home splash
+    if page == "Home":
+        display_sample_query_boxes(key_prefix="home_")
+
+    # Sub-apps
+    else:
+        # auto-run if launched from Home
+        if st.session_state["user_query"] and not st.session_state["auto_query_sent"]:
+            process_query(
+                {"ETF Search":"search_etf",
+                 "ETF Comparison":"comparison",
+                 "ETF Portfolio":"portfolio_projection"}[page],
+                st.session_state["user_query"]
+            )
+            st.session_state["auto_query_sent"] = True
+
+        task_map = {
+            "ETF Search": "search_etf",
+            "ETF Comparison": "comparison",
+            "ETF Portfolio": "portfolio_projection",
         }
-        </style>
-        """, unsafe_allow_html=True)
-    st.header("How can I assist you today?")
-    st.write("Find ETFs that align with your investment goals and sector interests, compare performance, estimate your portfolio, and get insights.")
-    display_sample_query_box("search_etf")
-    display_sample_query_box("comparison")
-    display_sample_query_box("portfolio_projection")
-
-# ─── MAIN ────────────────────────────────────────────────────────────────
-
-st.set_page_config(layout="wide")
-
-# initialize page and histories
-if "page" not in st.session_state:
-    st.session_state["page"] = "Home"
-for t in ["search_etf", "comparison", "portfolio_projection"]:
-    st.session_state.setdefault(f"all_chat_history_{t}", [])
-
-# sidebar navigation buttons
-st.sidebar.title("Navigation")
-if st.sidebar.button("🏠 Home"):
-    st.session_state["page"] = "Home"
-if st.sidebar.button("🔎 ETF Search"):
-    st.session_state["page"] = "ETF Search"
-if st.sidebar.button("⚖️ ETF Comparison"):
-    st.session_state["page"] = "ETF Comparison"
-if st.sidebar.button("💼 ETF Portfolio"):
-    st.session_state["page"] = "ETF Portfolio"
-
-# render page
-page = st.session_state["page"]
-st.title(page)
-
-if page == "Home":
-    display_explore_etfs_chat()
-
-elif page == "ETF Search":
-    display_sample_query_box("search_etf")
-    display_chat_history("search_etf")
-    q = st.chat_input("Search for ETFs…", key="in_search")
-    if q:
-        process_query("search_etf", q)
-
-elif page == "ETF Comparison":
-    display_sample_query_box("comparison")
-    display_chat_history("comparison")
-    q = st.chat_input("Compare ETFs…", key="in_comp")
-    if q:
-        process_query("comparison", q)
-
-elif page == "ETF Portfolio":
-    display_sample_query_box("portfolio_projection")
-    display_chat_history("portfolio_projection")
-    q = st.chat_input("Project portfolio…", key="in_port")
-    if q:
-        process_query("portfolio_projection", q)
+        task = task_map[page]
+
+        display_sample_query_boxes(key_prefix="sub_")
+        display_chat_history(task)
+
+        # chat input
+        prompt = {
+            "ETF Search": "Search for ETFs…",
+            "ETF Comparison": "Compare ETFs…",
+            "ETF Portfolio": "Project portfolio…",
+        }[page]
+
+        q = st.chat_input(prompt, key="inp_" + task)
+        if q:
+            process_query(task, q)
+
+if __name__ == "__main__":
+    main()