Update app.py

app.py

@@ -1,64 +1,120 @@
+import os
 import streamlit as st
+import uuid
+import pandas as pd
+import modules
+import torch
 from sentence_transformers import SentenceTransformer
 import faiss
 from transformers import AutoTokenizer, AutoModelForTokenClassification
-import modules
-import torch
 
-# ─── CACHED RESOURCES ────────────────────────────────────────────────────────
+# ─── CACHES ─────────────────────────────────────────────────────────────────
+
+@st.cache_data(show_spinner=False)
+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.
+    """
+    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)
+
+    # 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"})
+    )
+    return df_etf, df_analyst_report, available_tickers, df_annual_return_master
+
 @st.cache_resource(show_spinner=False)
-def load_search_index():
-    # 1. Load the SentenceTransformer
+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.
+    """
+    # 1) Ensure ETF data (with 'doc') is loaded
+    df_etf, _, _, _ = load_etf_data()
+
+    # 2) Load SentenceTransformer model
     repo_name = "hskwon7/paraphrase-MiniLM-L6-v2-ft-for-etf-semantic-search"
-    model = SentenceTransformer(repo_name)
+    model     = SentenceTransformer(repo_name)
 
-    # 2. Build your document texts & embeddings
-    df_etf, *_ = modules.load_etf_data()
-    df_etf["doc"] = df_etf.apply(modules.make_doc_text, axis=1)
+    # 3) Prepare list of tickers
     ticker_list = df_etf["Ticker"].tolist()
-    embs = model.encode(df_etf["doc"].tolist(), convert_to_numpy=True)  # no progress bar
 
-    # 3. Create FAISS index once
-    faiss.normalize_L2(embs)
-    index = faiss.IndexFlatIP(embs.shape[1])
-    index.add(embs)
+    # 4) Persist / load FAISS index
+    index_path = "etf_faiss.index"
+    if os.path.exists(index_path):
+        index = faiss.read_index(index_path)
+    else:
+        docs = df_etf["doc"].tolist()
+        embs = model.encode(docs, convert_to_numpy=True)
+        faiss.normalize_L2(embs)
+        index = faiss.IndexFlatIP(embs.shape[1])
+        index.add(embs)
+        faiss.write_index(index, index_path)
 
     return model, index, ticker_list
 
 @st.cache_resource(show_spinner=False)
 def load_ner_models():
-    # Load your two NER pipelines
-    repo1, repo2 = (
-        "hskwon7/distilbert-base-uncased-for-etf-ticker",
-        "hskwon7/albert-base-v2-for-etf-ticker",
-    )
+    """
+    Loads two NER models for ticker extraction, and builds the 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)
 
-    # valid tickers
-    df_etf, *_ = modules.load_etf_data()
-    valid = set(t.upper() for t in df_etf["Ticker"].unique())
+    df_etf, *_ = load_etf_data()
+    valid_ticker_set = set(t.upper() for t in df_etf["Ticker"].unique())
 
-    return (tok1, m1), (tok2, m2), valid
+    return (tok1, m1), (tok2, m2), valid_ticker_set
 
-# ─── IN YOUR APP ────────────────────────────────────────────────────────────
+# ─── INITIALIZE CACHED RESOURCES ─────────────────────────────────────────────
 
-# at the top of your script, instead of re-loading inline:
-s2_model, faiss_index, etf_list = load_search_index()
+# Load data & models once per session
+
+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()
 
-def semantic_search(query, top_k=100):
+# ─── CORE SEARCH & EXTRACTION ─────────────────────────────────────────────────
+
+def semantic_search(query: str, top_k: int = 100):
+    """
+    Returns a list of (ticker, score) tuples from FAISS semantic search.
+    """
     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 ensemble_ticker_extraction(query):
+
+def ensemble_ticker_extraction(query: str):
+    """
+    Runs two NER models in ensemble to extract possible ETF tickers.
+    """
     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()