Update app.py

app.py

@@ -9,8 +9,7 @@ Data sources, per the assignment constraints:
   * Embeddings      -> stored in THIS Space repo, with a dataset-repo fallback
 
 Runs on ZeroGPU: the generation model is placed on CUDA at module level and the
-single GPU-dependent function is decorated with @spaces.GPU, so a GPU is
-allocated only for the seconds it is actually needed.
+single GPU-dependent function is decorated with @spaces.GPU.
 """
 
 import os, json, re
@@ -31,17 +30,13 @@ DIV_COL      = "dividend_yield_10_year_pct"
 # LOAD ARTIFACTS
 # ---------------------------------------------------------------------------
 print("Loading artifacts...")
-
-# The config records which model built the index, so the query encoder can never
-# drift out of sync with the vectors it searches.
 cfg = json.load(open(hf_hub_download(DATASET_REPO, "embedding_config.json",
                                      repo_type="dataset")))
 EMB_MODEL_ID = cfg["model_id"]
 QUERY_PREFIX = cfg.get("query_prefix") or ""
 EMB_FILE     = f"stock_embeddings_{cfg['model_short']}.npy"
 
-# Embeddings: prefer the copy stored in this Space repo (no download at cold
-# start); fall back to the dataset repo if it is absent.
+# Prefer the copy stored in this Space repo; fall back to the dataset repo.
 if os.path.exists(EMB_FILE):
     embeddings = np.load(EMB_FILE)
     print(f"Embeddings loaded from Space repo: {EMB_FILE}")
@@ -49,7 +44,6 @@ else:
     embeddings = np.load(hf_hub_download(DATASET_REPO, EMB_FILE, repo_type="dataset"))
     print(f"Embeddings loaded from dataset repo: {EMB_FILE}")
 
-# Dataset and cached rationales: read directly from the HF dataset repo.
 df = pd.read_parquet(hf_hub_download(DATASET_REPO, "stock_metadata.parquet",
                                      repo_type="dataset"))
 QUICKSTART_CACHE = json.load(open(hf_hub_download(
@@ -61,26 +55,17 @@ assert len(embeddings) == len(df), "embeddings and metadata are misaligned"
 index = faiss.IndexFlatIP(embeddings.shape[1])
 index.add(embeddings.astype("float32"))
 
-# Embedding model pulled from the HF model repo. Kept on CPU: it runs outside
-# the @spaces.GPU function and encoding one short query takes milliseconds.
+# Embedding model from the HF model repo. Kept on CPU: it runs outside the
+# @spaces.GPU function and encoding one short query takes milliseconds.
 encoder = SentenceTransformer(EMB_MODEL_ID, device="cpu")
 print(f"Ready: {index.ntotal:,} stocks | encoder {EMB_MODEL_ID}")
 
-# ---------------------------------------------------------------------------
-# GENERATION MODEL — placed on CUDA at MODULE level, as ZeroGPU requires.
-# A PyTorch CUDA emulation mode is active outside @spaces.GPU functions, so this
-# works at startup; a real GPU is allocated only inside the decorated function.
-# Lazy loading is explicitly discouraged — CUDA transfers are optimised for
-# placement done during startup.
-#
-# SmolLM2-1.7B was selected on measured evidence: across 5 runs per model it was
-# both the fastest and the only candidate with zero fabricated figures. Qwen-1.5B
-# produced 9 fabrications, including an annual dividend income given as $800 on
-# one run and $9,000 on another where the correct value was $1,050.
-# ---------------------------------------------------------------------------
+# ZeroGPU requires models on CUDA at MODULE level. A CUDA emulation mode is
+# active outside @spaces.GPU functions, so this works at startup; a real GPU is
+# allocated only inside the decorated function.
 GEN_MODEL_ID = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
 gen_tok = AutoTokenizer.from_pretrained(GEN_MODEL_ID)
-gen_tok.padding_side = "left"                      # required for batched generation
+gen_tok.padding_side = "left"
 if gen_tok.pad_token is None:
     gen_tok.pad_token = gen_tok.eos_token
 gen_model = AutoModelForCausalLM.from_pretrained(
@@ -89,7 +74,10 @@ gen_model.eval()
 print(f"Generation model ready: {GEN_MODEL_ID}")
 
 # ---------------------------------------------------------------------------
-# QUESTIONNAIRE -> QUERY   (identical logic to the notebook)
+# QUESTIONNAIRE -> QUERY
+# Questions 1, 3, 4 and 6 accept MULTIPLE selections. Each selected option
+# contributes its phrase to the semantic query, and the hard filters use set
+# membership (isin) or a union of ranges rather than a single equality test.
 # ---------------------------------------------------------------------------
 RISK_TEXT = {
     "Low":    "low volatility, defensive, conservative risk, capital preservation",
@@ -101,100 +89,130 @@ STYLE_TEXT = {
     "Growth":   "pays no dividend, pure growth, high historical growth",
 }
 CAP_TEXT = {"Small": "small-cap company", "Medium": "mid-cap company",
-            "Large": "large-cap company, mega-cap company", "Any": ""}
+            "Large": "large-cap company, mega-cap company"}
 
 # Risk bands apply to BETA rather than the model's own risk_level label, which
 # the EDA showed sometimes contradicts the beta in the same record.
 RISK_BANDS = {"Low": (-0.5, 0.85), "Medium": (0.85, 1.30), "High": (1.30, 3.00)}
 
-# The dataset covers four regions across seven exchanges, so country answers map
-# to the nearest available region. Switzerland specifically is not representable.
 COUNTRY_TO_REGION = {
     "USA": "US", "UK": "Europe", "Switzerland": "Europe", "Germany": "Europe",
     "France": "Europe", "Europe": "Europe", "Israel": "Israel",
-    "Japan": "Asia", "China": "Asia", "Asia": "Asia", "Global": None,
+    "Japan": "Asia", "China": "Asia", "Asia": "Asia",
 }
-SECTORS = ["Any"] + sorted(df["sector"].unique().tolist())
-MARKETS = ["Global", "USA", "Europe", "Switzerland", "UK", "Israel", "Asia", "Japan"]
-
-def parse_markets(target_market):
-    if not target_market:
-        return None
-    items = ([m.strip() for m in target_market.split(",")]
-             if isinstance(target_market, str) else list(target_market))
-    if any(m in ("Global", "Any") for m in items):
+SECTORS = sorted(df["sector"].unique().tolist())
+MARKETS = ["USA", "Europe", "Switzerland", "UK", "Israel", "Asia", "Japan"]
+AMOUNTS = ["$1,000", "$5,000", "$10,000", "$25,000", "$50,000",
+           "$100,000", "$250,000", "$500,000"]
+
+def _as_list(x):
+    """Gradio multi-select returns a list; tolerate a bare string or None."""
+    if x is None:
+        return []
+    return list(x) if isinstance(x, (list, tuple)) else [x]
+
+def parse_markets(markets):
+    """Map selected countries to the four regions the dataset covers.
+    Empty selection means unrestricted."""
+    items = _as_list(markets)
+    if not items:
         return None
     regions = {COUNTRY_TO_REGION.get(m) for m in items}
     regions.discard(None)
     return sorted(regions) or None
 
-def build_query(risk, sector, market, style, cap_pref):
-    """Risk is stated twice: in a multi-clause query a single mention is diluted."""
-    parts = [RISK_TEXT.get(risk, ""), STYLE_TEXT.get(style, "")]
-    if CAP_TEXT.get(cap_pref):
-        parts.append(CAP_TEXT[cap_pref])
-    if sector and sector != "Any":
-        parts.append(f"{sector} company")
-    regions = parse_markets(market)
+def parse_amount(a):
+    """The dropdown shows formatted currency; strip it back to a number."""
+    if isinstance(a, str):
+        return float(re.sub(r"[^\d.]", "", a) or 0)
+    return float(a or 0)
+
+def build_query(risks, sectors, markets, style, caps):
+    """Every selected option contributes its phrase. The risk phrases are
+    repeated at the end because in a multi-clause query a single mention gets
+    diluted by the other attributes."""
+    risk_phrase = ". ".join(RISK_TEXT[r] for r in _as_list(risks) if r in RISK_TEXT)
+    parts = [risk_phrase, STYLE_TEXT.get(style, "")]
+    cap_phrase = ", ".join(CAP_TEXT[c] for c in _as_list(caps) if c in CAP_TEXT)
+    if cap_phrase:
+        parts.append(cap_phrase)
+    sec = _as_list(sectors)
+    if sec:
+        parts.append(" or ".join(f"{s} company" for s in sec))
+    regions = parse_markets(markets)
     if regions:
         parts.append("listed in " + " or ".join(regions))
-    parts.append(RISK_TEXT.get(risk, ""))
+    if risk_phrase:
+        parts.append(risk_phrase)
     return ". ".join([p for p in parts if p])
 
 def embed_query(text):
     return encoder.encode([QUERY_PREFIX + text],
                           normalize_embeddings=True).astype("float32")
 
-def recommend(risk_level, investment_amount, sector, target_market,
-              stock_type, market_cap_pref, top_k=3):
+def _beta_mask(frame, risks):
+    """Union of the selected risk bands. No selection means no restriction."""
+    lv = [r for r in _as_list(risks) if r in RISK_BANDS]
+    if not lv:
+        return pd.Series(True, index=frame.index)
+    m = pd.Series(False, index=frame.index)
+    for r in lv:
+        lo, hi = RISK_BANDS[r]
+        m |= (frame["beta"] >= lo) & (frame["beta"] < hi)
+    return m
+
+def recommend(risks, amount, sectors, markets, style, caps, top_k=3):
     """
     FILTER-THEN-RANK. Explicit constraints are enforced in pandas; semantic
     similarity ranks the eligible candidates. Risk is a hard filter because
     ablation showed embedding similarity could not separate Low from Medium
     (beta 1.205 vs 1.200 against a dataset mean of 1.21).
     """
-    query   = build_query(risk_level, sector, target_market, stock_type, market_cap_pref)
-    regions = parse_markets(target_market)
-    lo, hi  = RISK_BANDS[risk_level]
+    query   = build_query(risks, sectors, markets, style, caps)
+    regions = parse_markets(markets)
+    sec     = _as_list(sectors)
+    amount  = parse_amount(amount)
 
-    scores, idx = index.search(embed_query(query), k=min(top_k * 200, index.ntotal))
+    scores, idx = index.search(embed_query(query), k=min(top_k * 300, index.ntotal))
     res = df.iloc[idx[0]].copy()
     res["similarity"] = scores[0].round(4)
 
-    res = res[(res["beta"] >= lo) & (res["beta"] < hi)]
-    res = res[res["price_usd"] <= investment_amount]
-    if sector != "Any":
-        res = res[res["sector"] == sector]
+    res = res[_beta_mask(res, risks)]
+    res = res[res["price_usd"] <= amount]
+    if sec:
+        res = res[res["sector"].isin(sec)]
     if regions:
         res = res[res["market_region"].isin(regions)]
-    if stock_type == "Dividend":
+    if style == "Dividend":
         res = res[res[DIV_COL] > 0]
 
     note = ""
     if res.empty:
-        # Keep the SECTOR the user actively chose and widen the risk band
+        # Keep the SECTORS the user actively chose and widen the risk band
         # instead, reporting what was traded. Substituting a different industry
         # would ignore their stated intent.
         pool = df.copy()
-        if sector != "Any":
-            pool = pool[pool["sector"] == sector]
+        if sec:
+            pool = pool[pool["sector"].isin(sec)]
         if regions:
             pool = pool[pool["market_region"].isin(regions)]
-        pool = pool[pool["price_usd"] <= investment_amount]
-        if stock_type == "Dividend":
+        pool = pool[pool["price_usd"] <= amount]
+        if style == "Dividend":
             pool = pool[pool[DIV_COL] > 0]
         if not pool.empty:
-            pool = (pool.nsmallest(top_k, "beta") if risk_level == "Low"
-                    else pool.nlargest(top_k, "beta") if risk_level == "High"
+            lv = _as_list(risks)
+            pool = (pool.nsmallest(top_k, "beta") if lv == ["Low"]
+                    else pool.nlargest(top_k, "beta") if lv == ["High"]
                     else pool.iloc[(pool["beta"] - 1.075).abs().argsort()].head(top_k))
-            note = (f"⚠️ No **{sector}** stock falls in the **{risk_level}** risk band. "
-                    f"Showing the closest available {sector} stocks "
-                    f"(beta {pool['beta'].min():.2f}–{pool['beta'].max():.2f}). "
-                    f"Sector, market and budget constraints are still enforced.")
+            note = (f"⚠️ No stock in {', '.join(sec) or 'the selected sectors'} falls "
+                    f"in the {', '.join(lv) or 'selected'} risk band. Showing the "
+                    f"closest available (beta {pool['beta'].min():.2f}–"
+                    f"{pool['beta'].max():.2f}). Sector, market and budget "
+                    f"constraints are still enforced.")
             res = pool
         else:
             note = ("⚠️ No stock matches these constraints. "
-                    "Try widening your budget or markets.")
+                    "Try widening your budget, sectors or markets.")
             res = df.head(0)
     return res.head(top_k).reset_index(drop=True), note, query
 
@@ -221,16 +239,17 @@ def _tidy(s, max_words=30):
         s = s.rsplit(",", 1)[0]
     return s.rstrip(",;: ") + "."
 
-def _stock_prompt(risk, style, r):
+def _stock_prompt(risks, style, r):
     """Qualitative descriptors only — the model never sees a raw figure, so it
-    cannot perform arithmetic on one. This is a structural mitigation: given raw
-    numbers, the benchmark model fabricated 9 figures across 5 runs."""
+    cannot perform arithmetic on one. Given raw numbers during benchmarking, a
+    model fabricated 9 figures across 5 runs."""
     vol = "low" if r["beta"] < 0.85 else ("moderate" if r["beta"] <= 1.30 else "high")
     inc = ("no" if r[DIV_COL] == 0 else "low" if r[DIV_COL] < 1.5
            else "moderate" if r[DIV_COL] < 3.0 else "high")
     grw = ("slow" if r["cagr_10yr_pct"] < 5 else
            "steady" if r["cagr_10yr_pct"] < 12 else "fast")
-    return (f"Investor: {risk.lower()} risk tolerance, prefers {style.lower()} stocks.\n"
+    rl = "/".join(_as_list(risks)).lower() or "flexible"
+    return (f"Investor: {rl} risk tolerance, prefers {style.lower()} stocks.\n"
             f"This stock has {vol} volatility, {inc} dividend income, and {grw} "
             f"historical growth.\n"
             f"In one sentence of at most 20 words, explain what that means for this "
@@ -238,9 +257,7 @@ def _stock_prompt(risk, style, r):
 
 @spaces.GPU(duration=30)
 def _generate_sentences(prompts):
-    """The only function needing a real GPU, so the only one decorated.
-    duration=30 is generous for three short generations; shorter declared
-    durations receive higher queue priority."""
+    """The only function needing a real GPU, so the only one decorated."""
     enc = gen_tok(prompts, return_tensors="pt", padding=True).to("cuda")
     with torch.no_grad():
         out = gen_model.generate(**enc, max_new_tokens=110, do_sample=True,
@@ -250,22 +267,32 @@ def _generate_sentences(prompts):
     return gen_tok.batch_decode(out[:, enc["input_ids"].shape[1]:],
                                 skip_special_tokens=True)
 
-def generate_rationale(risk, amount, sector, market, style, recs):
+def generate_rationale(risks, amount, sectors, markets, style, recs):
     """One sentence per stock, generated in a single batched pass. All factual
     content — tickers and figures — is printed from the DataFrame, so it cannot
     be hallucinated."""
     prompts = [gen_tok.apply_chat_template(
         [{"role": "system", "content": STOCK_SYSTEM_PROMPT},
-         {"role": "user",   "content": _stock_prompt(risk, style, r)}],
+         {"role": "user",   "content": _stock_prompt(risks, style, r)}],
         tokenize=False, add_generation_prompt=True) for _, r in recs.iterrows()]
 
-    sents = _generate_sentences(prompts)
-
-    lines = [f"Based on your {risk.lower()}-risk profile and ${amount:,.0f} budget, "
-             f"here are {len(recs)} {sector} stocks from {market} matching your "
-             f"preferences.", ""]
+    try:
+        sents = _generate_sentences(prompts)
+    except Exception as e:
+        # A GPU failure must not take the whole recommendation down; the table
+        # above is still valid and useful on its own.
+        return (f"_Recommendations above are complete. The written explanation "
+                f"could not be generated: {type(e).__name__}._")
+
+    amt = parse_amount(amount)
+    rl  = "/".join(_as_list(risks)).lower() or "flexible"
+    sec = ", ".join(_as_list(sectors)) or "all sectors"
+    mkt = ", ".join(_as_list(markets)) or "all markets"
+
+    lines = [f"Based on your {rl}-risk profile and ${amt:,.0f} budget, here are "
+             f"{len(recs)} stocks in {sec} from {mkt} matching your preferences.", ""]
     for i, ((_, r), s) in enumerate(zip(recs.iterrows(), sents), 1):
-        sh = int(amount // r["price_usd"])
+        sh = int(amt // r["price_usd"]) if r["price_usd"] else 0
         lines.append(f"{i}. **{r['ticker']}** — beta {r['beta']:.2f}, "
                      f"{r[DIV_COL]:.2f}% yield, {r['cagr_10yr_pct']:.1f}% growth, "
                      f"${r['price_usd']:,.2f} per share ({sh} shares affordable). "
@@ -280,17 +307,15 @@ def generate_rationale(risk, amount, sector, market, style, recs):
 DISPLAY_COLS = ["ticker", "sector", "market_region", "cluster_name", "risk_level",
                 "beta", DIV_COL, "cagr_10yr_pct", "price_usd", "similarity"]
 
-def run_custom(risk, amount, sector, market, style, cap):
-    """Yields twice: the table appears immediately, the explanation follows."""
-    recs, note, query = recommend(risk, amount, sector, market, style, cap)
+def run_custom(risks, amount, sectors, markets, style, caps):
+    """Plain function, not a generator: @spaces.GPU does not compose well with
+    Gradio generator callbacks, and on GPU the whole round trip is ~2 seconds."""
+    recs, note, query = recommend(risks, amount, sectors, markets, style, caps)
     if recs.empty:
-        yield pd.DataFrame(), note, ""
-        return
+        return pd.DataFrame(), note, ""
     cols = [c for c in DISPLAY_COLS if c in recs.columns]
-    yield recs[cols], (note + "\n\n⏳ Generating your personalised explanation…").strip(), \
-          f"*Semantic query:* `{query}`"
-    text = generate_rationale(risk, amount, sector, market, style, recs)
-    yield recs[cols], text, f"*Semantic query:* `{query}`"
+    text = generate_rationale(risks, amount, sectors, markets, style, recs)
+    return recs[cols], (note + "\n\n" + text).strip(), f"*Semantic query:* `{query}`"
 
 def run_quickstart(name):
     """Served from the pre-generated cache — instant, uses no GPU quota."""
@@ -327,18 +352,23 @@ with gr.Blocks(title="StockMatch AI", theme=gr.themes.Soft()) as demo:
         b2 = gr.Button("👔 Balanced Professional", variant="secondary")
         b3 = gr.Button("🚀 Young Growth Seeker", variant="secondary")
 
-    gr.Markdown("### 📋 Or build your own profile")
+    gr.Markdown("### 📋 Or build your own profile — questions 1, 3, 4 and 6 accept "
+                "multiple answers")
     with gr.Row():
         with gr.Column():
-            q1 = gr.Radio(["Low", "Medium", "High"], value="Low", label="1. Risk level")
-            q2 = gr.Number(value=25000, label="2. Investment amount (USD)")
-            q3 = gr.Dropdown(SECTORS, value="Utilities", label="3. Sector")
+            q1 = gr.CheckboxGroup(["Low", "Medium", "High"], value=["Low"],
+                                  label="1. Risk level (select one or more)")
+            q2 = gr.Dropdown(AMOUNTS, value="$25,000",
+                             label="2. Investment amount (USD)")
+            q3 = gr.Dropdown(SECTORS, value=["Utilities"], multiselect=True,
+                             label="3. Sector (select one or more, or none for all)")
         with gr.Column():
-            q4 = gr.Dropdown(MARKETS, value="USA", label="4. Target market")
+            q4 = gr.Dropdown(MARKETS, value=["USA"], multiselect=True,
+                             label="4. Target market (select one or more, or none for all)")
             q5 = gr.Radio(["Dividend", "Growth"], value="Dividend",
                           label="5. Stock type preference")
-            q6 = gr.Radio(["Small", "Medium", "Large", "Any"], value="Large",
-                          label="6. Company size preference")
+            q6 = gr.CheckboxGroup(["Small", "Medium", "Large"], value=["Large"],
+                                  label="6. Company size (select one or more)")
     go = gr.Button("🔍 Find my stocks", variant="primary", size="lg")
 
     gr.Markdown("### Results")