Update app.py

app.py

@@ -8,8 +8,9 @@ Data sources, per the assignment constraints:
   * Embedding model -> read from the HF MODEL repo   (BAAI/bge-small-en-v1.5)
   * 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.
+Runs on free-tier CPU. The three Quick Starters are served from a pre-generated
+cache so the common path is instant; the language model is loaded lazily and
+only when a custom questionnaire is submitted.
 """
 
 import os, json, re
@@ -35,6 +36,7 @@ EMB_MODEL_ID = cfg["model_id"]
 QUERY_PREFIX = cfg.get("query_prefix") or ""
 EMB_FILE     = f"stock_embeddings_{cfg['model_short']}.npy"
 
+# 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}")
@@ -53,21 +55,25 @@ assert len(embeddings) == len(df), "embeddings and metadata are misaligned"
 index = faiss.IndexFlatIP(embeddings.shape[1])
 index.add(embeddings.astype("float32"))
 
-# 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}")
 
-# ZeroGPU requires models on CUDA at MODULE level.
+# The generation model is heavy on CPU, so it is loaded on FIRST USE only.
+# Quick Starters never trigger it — they are served from the cache.
 GEN_MODEL_ID = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
-gen_tok = AutoTokenizer.from_pretrained(GEN_MODEL_ID)
-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(
-    GEN_MODEL_ID, dtype=torch.float16).to("cuda")
-gen_model.eval()
-print(f"Generation model ready: {GEN_MODEL_ID}")
+_gen = {}
+
+def _load_generator():
+    if "model" not in _gen:
+        print(f"Lazy-loading {GEN_MODEL_ID}...")
+        tok = AutoTokenizer.from_pretrained(GEN_MODEL_ID)
+        tok.padding_side = "left"              # required for batched generation
+        if tok.pad_token is None:
+            tok.pad_token = tok.eos_token
+        m = AutoModelForCausalLM.from_pretrained(GEN_MODEL_ID, dtype=torch.float32)
+        m.eval()
+        _gen.update(tok=tok, model=m)
+    return _gen["tok"], _gen["model"]
 
 # ---------------------------------------------------------------------------
 # QUESTIONNAIRE -> QUERY
@@ -87,8 +93,12 @@ STYLE_TEXT = {
 CAP_TEXT = {"Small": "small-cap company", "Medium": "mid-cap company",
             "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 built from seven exchanges, so country choices
+# map upward: Japan and China -> Asia, Switzerland and the UK -> Europe.
 COUNTRY_TO_REGION = {
     "USA": "US", "UK": "Europe", "Switzerland": "Europe", "Germany": "Europe",
     "France": "Europe", "Europe": "Europe", "Israel": "Israel",
@@ -100,6 +110,7 @@ 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]
@@ -226,6 +237,7 @@ _DANGLING = {"a","an","the","and","or","but","though","with","for","in","to","of
              "its","their","may","can","will","while","as","that","this","from","is"}
 
 def _tidy(s, max_words=30):
+    """Keep the first sentence, cap its length, never end mid-clause."""
     s = " ".join(s.strip().split()).split(". ")[0].rstrip(". ")
     w = s.split()
     if len(w) > max_words:
@@ -236,7 +248,8 @@ def _tidy(s, max_words=30):
 
 def _stock_prompt(risks, style, r):
     """Qualitative descriptors only — the model never sees a raw figure, so it
-    cannot perform arithmetic on one."""
+    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")
@@ -250,43 +263,20 @@ def _stock_prompt(risks, style, r):
             f"In one sentence of at most 20 words, explain what that means for this "
             f"investor and note one caveat.")
 
-@spaces.GPU(duration=30)
 def _generate_sentences(prompts):
-    """The only function needing a real GPU, so the only one decorated.
-
-    Errors are caught HERE and returned as data rather than raised. ZeroGPU
-    executes this in a separate worker process, and an exception crossing that
-    boundary arrives as a bare RuntimeError with its message stripped — which
-    makes diagnosis impossible. Returning the traceback preserves it.
-    """
-    import traceback
-    try:
-        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,
-                                     temperature=0.4, top_p=0.9,
-                                     repetition_penalty=1.1,
-                                     pad_token_id=gen_tok.pad_token_id)
-        sents = gen_tok.batch_decode(out[:, enc["input_ids"].shape[1]:],
-                                     skip_special_tokens=True)
-        return {"ok": True, "sentences": sents}
-    except Exception:
-        return {"ok": False, "traceback": traceback.format_exc()}
-
-
-def _generate_sentences(prompts):
-    """Runs on CPU. Errors are returned as data rather than raised so the
-    recommendation table survives a generation failure."""
+    """Runs on CPU. Errors are returned as data rather than raised, so a
+    generation failure never takes down the recommendation table."""
     import traceback
     try:
-        enc = gen_tok(prompts, return_tensors="pt", padding=True)
+        tok, model = _load_generator()
+        enc = tok(prompts, return_tensors="pt", padding=True)
         with torch.no_grad():
-            out = gen_model.generate(**enc, max_new_tokens=110, do_sample=True,
-                                     temperature=0.4, top_p=0.9,
-                                     repetition_penalty=1.1,
-                                     pad_token_id=gen_tok.pad_token_id)
-        sents = gen_tok.batch_decode(out[:, enc["input_ids"].shape[1]:],
-                                     skip_special_tokens=True)
+            out = model.generate(**enc, max_new_tokens=110, do_sample=True,
+                                 temperature=0.4, top_p=0.9,
+                                 repetition_penalty=1.1,
+                                 pad_token_id=tok.pad_token_id)
+        sents = tok.batch_decode(out[:, enc["input_ids"].shape[1]:],
+                                 skip_special_tokens=True)
         return {"ok": True, "sentences": sents}
     except Exception:
         return {"ok": False, "traceback": traceback.format_exc()}
@@ -295,7 +285,8 @@ 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(
+    tok, _ = _load_generator()
+    prompts = [tok.apply_chat_template(
         [{"role": "system", "content": STOCK_SYSTEM_PROMPT},
          {"role": "user",   "content": _stock_prompt(risks, style, r)}],
         tokenize=False, add_generation_prompt=True) for _, r in recs.iterrows()]
@@ -327,8 +318,8 @@ def generate_rationale(risks, amount, sectors, markets, style, recs):
 
 # ---------------------------------------------------------------------------
 # RESULTS TABLE FORMATTING
-# Database column names are renamed to plain English and numbers are rounded,
-# so the table reads as a product rather than a database dump.
+# Database column names are renamed to plain English and numbers rounded, so the
+# table reads as a product rather than a database dump.
 # ---------------------------------------------------------------------------
 DISPLAY_COLS = ["ticker", "sector", "market_region", "cluster_name", "risk_level",
                 "beta", DIV_COL, "cagr_10yr_pct", "price_usd", "similarity"]
@@ -348,27 +339,33 @@ COLUMN_LABELS = {
 
 def _prettify(frame):
     out = frame[[c for c in DISPLAY_COLS if c in frame.columns]].copy()
+    # FAISS returns float32; rounding without casting to float64 leaves
+    # artefacts such as 75.30000305175781 on screen.
     if "similarity" in out:
-        out["similarity"] = (out["similarity"] * 100).round(1)   # 0.804 -> 80.4
-    for c in ["beta", DIV_COL, "cagr_10yr_pct"]:
+        out["similarity"] = (out["similarity"].astype("float64") * 100).round(1)
+    for c in ["beta", DIV_COL, "cagr_10yr_pct", "price_usd"]:
         if c in out:
-            out[c] = out[c].round(2)
-    if "price_usd" in out:
-        out["price_usd"] = out["price_usd"].round(2)
+            out[c] = out[c].astype("float64").round(2)
     return out.rename(columns=COLUMN_LABELS)
 
 # ---------------------------------------------------------------------------
 # GRADIO CALLBACKS
 # ---------------------------------------------------------------------------
 def run_custom(risks, amount, sectors, markets, style, caps):
+    """Yields twice: the table appears immediately, then the explanation, so a
+    CPU generation of roughly a minute does not look like a frozen page."""
     recs, note, _query = recommend(risks, amount, sectors, markets, style, caps)
     if recs.empty:
-        return pd.DataFrame(), note
+        yield pd.DataFrame(), note
+        return
+    pretty = _prettify(recs)
+    yield pretty, (note + "\n⏳ Writing your personalised explanation… "
+                          "(about a minute on free CPU)").strip()
     text = generate_rationale(risks, amount, sectors, markets, style, recs)
-    return _prettify(recs), (note + "\n" + text).strip()
+    yield pretty, (note + "\n" + text).strip()
 
 def run_quickstart(name):
-    """Served from the pre-generated cache — instant, uses no GPU quota."""
+    """Served from the pre-generated cache — instant, no model load."""
     entry = QUICKSTART_CACHE[name]
     tbl = pd.DataFrame(entry["table"])
     for c in ["beta", DIV_COL, "cagr_10yr_pct", "price_usd", "similarity",
@@ -418,10 +415,6 @@ with gr.Blocks(title="StockMatch AI", theme=gr.themes.Soft()) as demo:
                                   label="6. Company size")
     go = gr.Button("🔍 Find my stocks", variant="primary", size="lg")
 
-    # Diagnostic only — remove before submitting.
-    gpu_btn = gr.Button("🔧 Test GPU", variant="secondary", size="sm")
-    gpu_out = gr.Markdown()
-
     gr.Markdown("### Results")
     table = gr.Dataframe(label="Your matches", interactive=False, wrap=True)
     text  = gr.Markdown()
@@ -430,6 +423,5 @@ with gr.Blocks(title="StockMatch AI", theme=gr.themes.Soft()) as demo:
     b1.click(lambda: run_quickstart("Cautious Retiree"),      None, [table, text])
     b2.click(lambda: run_quickstart("Balanced Professional"), None, [table, text])
     b3.click(lambda: run_quickstart("Young Growth Seeker"),   None, [table, text])
-    gpu_btn.click(_gpu_selftest, None, gpu_out)
 
 demo.launch()