Update app.py

app.py

@@ -1,130 +1,209 @@
-# app.py
-
-"""
-Streamlit application for Image-to-Story demo with history sidebar.
-Allows demo/upload image, generates a caption, a trimmed story,
-and plays back as MP3 via gTTS. Keeps history of all runs.
-"""
 import streamlit as st
-from PIL import Image
-import warnings
-from modules import (
-    load_captioner, load_story_gen,
-    generate_caption, generate_story_simple,
-    generate_audio
-)
-import io
-
-warnings.filterwarnings("ignore", category=DeprecationWarning)
-
-# Reset state when switching image source
-def reset_state():
-    for key in ["caption", "story", "audio_bytes", "audio_mime", "selected_index"]:
-        if key in st.session_state:
-            del st.session_state[key]
-
-def main():
-    st.set_page_config(layout="wide")
-    st.title("🎨 Magic Picture Story Time!")
-    st.write("Pick or upload a picture, and watch it turn into a fun story with voice! Ready for a magical tale?")
-
-    # --- Sidebar: History ---
-    st.sidebar.header("History")
-    if "history" not in st.session_state:
-        st.session_state.history = []           # list of dicts
-    if "selected_index" not in st.session_state:
-        st.session_state.selected_index = None
-
-    # Render thumbnails & select buttons
-    for idx, entry in enumerate(st.session_state.history):
-        with st.sidebar.container():
-            st.sidebar.image(entry["image_bytes"], width=100)
-            if st.sidebar.button(f"View #{idx+1}", key=f"view_{idx}"):
-                st.session_state.selected_index = idx
-
-    # Sidebar clear-all button
-    if st.sidebar.button("Clear History"):
-        st.session_state.history = []
-        st.session_state.selected_index = None
-
-    # --- Main panel: image selection ---
-    source = st.radio("Image source:",
-                      ("Upload my own image", "Use demo image"),
-                      on_change=reset_state)
-
-    # Load pipelines once
-    if "models_loaded" not in st.session_state:
-        with st.spinner("Loading models…"):
-            st.session_state.captioner = load_captioner()
-            st.session_state.story_gen = load_story_gen()
-        st.session_state.models_loaded = True
-
-    captioner = st.session_state.captioner
-    story_gen = st.session_state.story_gen
-
-    # If user clicked a history entry, load it
-    sel = st.session_state.selected_index
-    if sel is not None:
-        entry = st.session_state.history[sel]
-        img = Image.open(io.BytesIO(entry["image_bytes"])).convert("RGB")
-        st.image(img, use_container_width=True)
-        st.markdown(f"**Caption:** {entry['caption']}")
-        st.markdown(f"**Story:** {entry['story']}")
-        if st.button("🔊 Play Story Audio"):
-            st.audio(data=entry["audio_bytes"], format=entry["audio_mime"])
-        return
-
-    # Otherwise, handle a fresh upload/demo
-    if source == "Use demo image":
-        img = Image.open("test_kids_playing.jpg").convert("RGB")
-        # grab raw bytes for history
-        buf = io.BytesIO()
-        img.save(buf, format="JPEG")
-        img_bytes = buf.getvalue()
-    else:
-        uploaded = st.file_uploader("Upload an image",
-                                    type=["png", "jpg", "jpeg"])
-        if not uploaded:
-            return
-        img = Image.open(uploaded).convert("RGB")
-        img_bytes = uploaded.getvalue()
-
-    st.image(img, use_container_width=True)
-
-    # Step 1: Caption
-    if "caption" not in st.session_state:
-        with st.spinner("Captioning image…"):
-            st.session_state.caption = generate_caption(captioner, img)
-    st.markdown(f"**Caption:** {st.session_state.caption}")
-
-    # Step 2: Story
-    if "story" not in st.session_state:
-        with st.spinner("Creating story…"):
-            st.session_state.story = generate_story_simple(
-                story_gen, st.session_state.caption, 50, 100
+import openai
+import uuid
+import modules
+import pandas as pd
+import torch
+
+from sentence_transformers import SentenceTransformer
+import faiss
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+
+# ─── LOAD DATA & MODELS (unchanged) ─────────────────────────────────────────
+df_etf, df_analyst_report, available_tickers, df_annual_return_master = modules.load_etf_data()
+
+repo_name = "hskwon7/paraphrase-MiniLM-L6-v2-ft-for-etf-semantic-search"
+s2_model = SentenceTransformer(repo_name)
+df_etf["doc"] = df_etf.apply(modules.make_doc_text, axis=1)
+etf_list  = df_etf["ticker"].tolist()
+doc_embs  = s2_model.encode(df_etf["doc"].tolist(), convert_to_numpy=True, show_progress_bar=True)
+faiss.normalize_L2(doc_embs)
+index = faiss.IndexFlatIP(doc_embs.shape[1])
+index.add(doc_embs)
+
+def semantic_search(query, top_k=100):
+    q_emb = s2_model.encode([query], convert_to_numpy=True)
+    faiss.normalize_L2(q_emb)
+    D, I = index.search(q_emb, top_k)
+    return [(etf_list[idx], float(D[0][i])) for i, idx in enumerate(I[0])]
+
+# NER ensemble remains the same
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+repo1 = "hskwon7/distilbert-base-uncased-for-etf-ticker"
+repo2 = "hskwon7/albert-base-v2-for-etf-ticker"
+tokenizer1 = AutoTokenizer.from_pretrained(repo1)
+model1     = AutoModelForTokenClassification.from_pretrained(repo1).to(device).eval()
+tokenizer2 = AutoTokenizer.from_pretrained(repo2)
+model2     = AutoModelForTokenClassification.from_pretrained(repo2).to(device).eval()
+valid_ticker_set = set(t.upper() for t in df_etf['ticker'].unique().tolist())
+
+def ensemble_ticker_extraction(query):
+    preds = set()
+    for tok, mdl in ((tokenizer1, model1),(tokenizer2, model2)):
+        enc = tok(query, return_tensors="pt").to(device)
+        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]
+        preds |= modules.extract_valid_tickers(toks, labs, tok, valid_ticker_set)
+    return preds
+
+# ─── HELPERS ──────────────────────────────────────────────────────────────
+
+def display_sample_query_box(task: str):
+    """
+    Renders one sample-query box for the given task
+    and adds a 'Try this app' button that switches the sidebar page.
+    """
+    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']
+        },
+        "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"]
+        },
+        "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"]
+        }
+    }
+    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()
+
+def display_chat_history(task: str):
+    """
+    Shows only the chat history for a given task.
+    """
+    hist = st.session_state.get(f"all_chat_history_{task}", [])
+    for entry in hist:
+        if task == "search_etf":
+            st.chat_message("assistant").write(entry["response"])
+            modules.display_matching_etfs(entry["df"])
+        elif task == "comparison":
+            st.chat_message("assistant").write(entry["response"])
+            st.plotly_chart(entry["fig"], use_container_width=True)
+            st.dataframe(entry["df"], hide_index=True)
+        elif task == "portfolio_projection":
+            st.chat_message("assistant").write(entry["response"])
+            st.plotly_chart(entry["fig"], use_container_width=True)
+
+def process_query(task: str, query: str):
+    """
+    Core logic for each sub-app.
+    """
+    # make sure top_k / top_n are always available
+    top_k, top_n = 100, 30
+
+    if task == 'search_etf':
+        with st.spinner("Searching ETFs..."):
+            fetched = semantic_search(query, top_k=top_k)
+            df_out  = modules.get_etf_recommendations_from_list(fetched,
+                        modules.get_cols_to_display(), df_etf, top_n=top_n)
+        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
+            )
+        )
+
+    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
+                )
+            else:
+                df_out = modules.get_etf_recommendations_from_list(
+                    [(t, None) for t in tickers],
+                    modules.get_cols_to_display(), df_etf, top_n=2
+                )
+                fig = modules.compare_etfs_interactive(*tickers)
+                response = f"Compared {tickers[0]} vs. {tickers[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
+            )
+        )
+
+    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)
+            fig     = modules.portfolio_interactive_chart(df_port)
+            response = "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
             )
-    st.markdown(f"**Story:** {st.session_state.story}")
-
-    # Step 3: Audio
-    if "audio_bytes" not in st.session_state:
-        with st.spinner("Generating audio…"):
-            audio_bytes, mime = generate_audio(st.session_state.story)
-            st.session_state.audio_bytes = audio_bytes
-            st.session_state.audio_mime  = mime
-
-    if st.button("🔊 Play Story Audio"):
-        st.audio(data=st.session_state.audio_bytes,
-                 format=st.session_state.audio_mime)
-
-    # Step 4: Append to history (only once per new run)
-    if not st.session_state.history or st.session_state.history[-1]["image_bytes"] != img_bytes:
-        st.session_state.history.append({
-            "image_bytes": img_bytes,
-            "caption":     st.session_state.caption,
-            "story":       st.session_state.story,
-            "audio_bytes": st.session_state.audio_bytes,
-            "audio_mime":  st.session_state.audio_mime
-        })
-
-if __name__ == "__main__":
-    main()
+        )
+
+# ─── MAIN ────────────────────────────────────────────────────────────────
+st.set_page_config(layout="wide")
+if "page" not in st.session_state:
+    st.session_state["page"] = "ETF Search"  # default
+
+# initialize histories
+for t in ["search_etf","comparison","portfolio_projection"]:
+    st.session_state.setdefault(f"all_chat_history_{t}", [])
+
+# sidebar navigation
+st.sidebar.title("ETF Assistant")
+st.sidebar.radio("Go to…", ["ETF Search","ETF Comparison","ETF Portfolio"], key="page")
+
+# dispatch
+page = st.session_state["page"]
+st.title(page)
+
+if 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)