Update app.py

app.py

@@ -12,64 +12,54 @@ from transformers import AutoTokenizer, AutoModelForTokenClassification
 
 @st.cache_data(show_spinner=False)
 def load_etf_data():
-    """
-    Load ETF data with a persistent 'doc' column.
-    - 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"
-
     if os.path.exists(enriched_path):
         df_info = pd.read_csv(enriched_path)
     else:
         df_info = pd.read_csv(raw_path).rename(columns={"ticker": "Ticker"})
         df_info["doc"] = df_info.apply(modules.make_doc_text, axis=1)
         df_info.to_csv(enriched_path, index=False)
-
     df_etf, available_tickers = modules.set_etf_data(df_info)
     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():
-    """
-    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()
-    repo_name  = "hskwon7/paraphrase-MiniLM-L6-v2-ft-for-etf-semantic-search"
-    model      = SentenceTransformer(repo_name)
+    model      = SentenceTransformer(
+        "hskwon7/paraphrase-MiniLM-L6-v2-ft-for-etf-semantic-search"
+    )
     ticker_list = df_etf["Ticker"].tolist()
 
-    index_path = "etf_faiss.index"
-    if os.path.exists(index_path):
-        index = faiss.read_index(index_path)
+    idx_path = "etf_faiss.index"
+    if os.path.exists(idx_path):
+        index = faiss.read_index(idx_path)
     else:
-        docs = df_etf["doc"].tolist()
-        embs = model.encode(docs, convert_to_numpy=True)
+        embs = model.encode(df_etf["doc"].tolist(), convert_to_numpy=True)
         faiss.normalize_L2(embs)
         index = faiss.IndexFlatIP(embs.shape[1])
         index.add(embs)
-        faiss.write_index(index, index_path)
+        faiss.write_index(index, idx_path)
 
     return model, index, ticker_list
 
 @st.cache_resource(show_spinner=False)
 def load_ner_models():
-    """
-    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, m1 = AutoTokenizer.from_pretrained(repo1), AutoModelForTokenClassification.from_pretrained(repo1)
-    tok2, m2 = AutoTokenizer.from_pretrained(repo2), AutoModelForTokenClassification.from_pretrained(repo2)
-
+    tok1, m1 = (
+        AutoTokenizer.from_pretrained("hskwon7/distilbert-base-uncased-for-etf-ticker"),
+        AutoModelForTokenClassification.from_pretrained("hskwon7/distilbert-base-uncased-for-etf-ticker")
+    )
+    tok2, m2 = (
+        AutoTokenizer.from_pretrained("hskwon7/albert-base-v2-for-etf-ticker"),
+        AutoModelForTokenClassification.from_pretrained("hskwon7/albert-base-v2-for-etf-ticker")
+    )
     df_etf, *_ = load_etf_data()
-    valid_ticker_set = set(t.upper() for t in df_etf["Ticker"].unique())
-
+    valid_ticker_set = set(df_etf["Ticker"].str.upper())
     return (tok1, m1), (tok2, m2), valid_ticker_set
 
 # ─── INITIALIZE ─────────────────────────────────────────────────────────────
@@ -78,18 +68,18 @@ df_etf, df_analyst_report, available_tickers, df_annual_return_master = load_etf
 s2_model, faiss_index, etf_list = build_search_resources()
 (tok1, m1), (tok2, m2), valid_ticker_set = load_ner_models()
 
-# ─── CORE SEARCH & EXTRACTION ───────────────────────────────────────────────
+# ─── CORE ROUTINES ──────────────────────────────────────────────────────────
 
-def semantic_search(query: str, top_k: int = 100):
-    q_emb = s2_model.encode([query], convert_to_numpy=True)
-    faiss.normalize_L2(q_emb)
-    D, I = faiss_index.search(q_emb, top_k)
-    return [(etf_list[idx], float(D[0][i])) for i, idx in enumerate(I[0])]
+def semantic_search(q: str, top_k: int=100):
+    emb = s2_model.encode([q], convert_to_numpy=True)
+    faiss.normalize_L2(emb)
+    D, I = faiss_index.search(emb, top_k)
+    return [(etf_list[i], float(D[0][j])) for j,i in enumerate(I[0])]
 
-def ensemble_ticker_extraction(query: str):
+def ensemble_ticker_extraction(q: str):
     preds = set()
-    for tok, mdl in ((tok1, m1), (tok2, m2)):
-        enc = tok(query, return_tensors="pt")
+    for tok, mdl in ((tok1,m1),(tok2,m2)):
+        enc = tok(q, return_tensors="pt")
         with torch.no_grad():
             logits = mdl(**enc).logits
         ids  = logits.argmax(dim=-1)[0].tolist()
@@ -98,89 +88,78 @@ def ensemble_ticker_extraction(query: str):
         preds |= modules.extract_valid_tickers(toks, labs, tok, valid_ticker_set)
     return preds
 
-# ─── UI HELPERS ──────────────────────────────────────────────────────────────
+# ─── UI HELPERS ─────────────────────────────────────────────────────────────
 
 def display_sample_query_boxes(key_prefix=""):
     sample_queries = {
         "search_etf": {
             "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?']
+            "query": [
+                'High-dividend ETFs in the tech sector',
+                'Precious metals ETFs with low expense ratio',
+                'Large growth ETFs with high returns'
+            ]
         },
         "comparison": {
             "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."
+            "query": [
+                "I'd like to compare performance of QQQ with GLD.",
+                "Compare SPY and VOO.",
+                "SCHD vs. VTI"
+            ]
         },
         "portfolio_projection": {
             "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."
+            "description": "Project a portfolio with multiple ETFs over 30 years.",
+            "query": [
+                "I want to invest in SPY, QQQ, SCHD, and IAU.",
+                "Portfolio projection for VTI, XLF, and XLY."
+            ]
         },
     }
 
     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)
+    title_h, desc_h, query_h = "40px", "60px", "60px"
 
     for idx, (key, details) in enumerate(sample_queries.items()):
         with cols[idx]:
-            box_html = f"""
+            st.markdown(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;
+              width:100%; height:300px; border:1px solid #ddd;
+              border-radius:10px; padding:15px; margin:auto;
+              display:flex; flex-direction:column; justify-content:space-between;
+              box-shadow:2px 2px 8px rgba(0,0,0,0.1);
             ">
-              <div style="height: {title_h}; text-align: center;">
-                <p style="margin:5px 0; color:#2c3e50; font-size:16px;">
-                  <b>{details['title']}</b>
-                </p>
+              <div style="height:{title_h}; text-align:center;">
+                <b style="font-size:16px; color:#2c3e50;">
+                  {details['title']}
+                </b>
               </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 style="height:{desc_h}; text-align:center; color:#7f8c8d; font-size:14px; overflow:auto;">
+                {details['description']}
               </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 style="height:{query_h}; text-align:center; color:#34495e; font-size:13px; font-style:italic; overflow:auto;">
+                {'<br>'.join(f'“{q}”' for q in details['query'])}
               </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()
+            """, unsafe_allow_html=True)
+
+            # center the button directly under the box
+            st.markdown("<div style='text-align:center; margin-top:10px;'>", unsafe_allow_html=True)
+            if st.button("Try this app", key=key_prefix+key):
+                page_map = {
+                    "search_etf": "ETF Search",
+                    "comparison": "ETF Comparison",
+                    "portfolio_projection": "ETF Portfolio"
+                }
+                st.session_state["page"] = page_map[key]
+                st.experimental_rerun()
+            st.markdown("</div>", unsafe_allow_html=True)
 
 def display_chat_history(task: str):
-    hist = st.session_state.get(f"all_chat_history_{task}", [])
-    for entry in hist:
+    for entry in st.session_state.get(f"all_chat_history_{task}", []):
         st.chat_message("assistant").write(entry["response"])
         if entry.get("fig"):
             st.plotly_chart(entry["fig"], use_container_width=True)
@@ -189,8 +168,7 @@ 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)
             df_out  = modules.get_etf_recommendations_from_list(
@@ -201,14 +179,14 @@ def process_query(task: str, query: str):
             modules.form_d_chat_history(str(uuid.uuid4()), msg, task, df=df_out)
         )
 
-    elif task == "comparison":
+    elif task=="comparison":
         with st.spinner("Running comparison..."):
             tk = ensemble_ticker_extraction(query)
-            if len(tk) != 2:
+            if len(tk)!=2:
                 resp, fig, df_out = "Please specify exactly two tickers.", None, None
             else:
                 df_out = modules.get_etf_recommendations_from_list(
-                    [(t, None) for t in tk],
+                    [(t,None) for t in tk],
                     modules.get_cols_to_display(), df_etf, top_n=2
                 )
                 fig  = modules.compare_etfs_interactive(*tk)
@@ -217,7 +195,7 @@ def process_query(task: str, query: str):
             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)
             df_port = modules.run_portfolio_analysis(fetched, df_etf, df_annual_return_master)
@@ -227,73 +205,46 @@ def process_query(task: str, query: str):
             modules.form_d_chat_history(str(uuid.uuid4()), resp, task, fig=fig)
         )
 
+# ─── MAIN ────────────────────────────────────────────────────────────────
+
 def main():
     st.set_page_config(layout="wide")
-
-    # init state
+    # init
     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
+        st.session_state["page"]="Home"
     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"
+    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)
+    st.title(page if page!="Home" else "ETF Assistant")
 
-    # Home splash
-    if page == "Home":
+    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",
-        }
-        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…",
+        task = {
+            "ETF Search":"search_etf",
+            "ETF Comparison":"comparison",
+            "ETF Portfolio":"portfolio_projection"
         }[page]
-
-        q = st.chat_input(prompt, key="inp_" + task)
+        display_chat_history(task)
+        q = st.chat_input({
+            "ETF Search":"Search for ETFs…",
+            "ETF Comparison":"Compare ETFs…",
+            "ETF Portfolio":"Project portfolio…"
+        }[page], key=task)
         if q:
             process_query(task, q)
 
-if __name__ == "__main__":
+if __name__=="__main__":
     main()