Update app.py

app.py

@@ -4,13 +4,13 @@ StockMatch AI — Gradio application (Part 5)
 QUESTIONNAIRE -> EMBEDDING -> FAISS RETRIEVAL -> FILTERS -> GENERATED RATIONALE
 
 Data sources, per the assignment constraints:
-  * Dataset      -> read directly from the HF DATASET repo (Kogann/stockmatch-synthetic)
-  * Embedding model -> read directly from the HF MODEL repo (BAAI/bge-small-en-v1.5)
-  * Embeddings   -> stored in THIS Space repo, with a dataset-repo fallback
+  * Dataset         -> read from the HF DATASET repo (Kogann/stockmatch-synthetic)
+  * 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
 
-The language model loads LAZILY, only when a custom questionnaire is submitted.
-The three Quick Starters are served from a pre-generated cache, so the common
-path never pays the model-load cost on free-tier CPU.
+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.
 """
 
 import os, json, re
@@ -18,8 +18,11 @@ import numpy as np
 import pandas as pd
 import gradio as gr
 import faiss
+import spaces
+import torch
 from huggingface_hub import hf_hub_download
 from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 DATASET_REPO = "Kogann/stockmatch-synthetic"
 DIV_COL      = "dividend_yield_10_year_pct"
@@ -29,7 +32,7 @@ DIV_COL      = "dividend_yield_10_year_pct"
 # ---------------------------------------------------------------------------
 print("Loading artifacts...")
 
-# Config records which model built the index, so the query encoder can never
+# 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")))
@@ -58,27 +61,32 @@ assert len(embeddings) == len(df), "embeddings and metadata are misaligned"
 index = faiss.IndexFlatIP(embeddings.shape[1])
 index.add(embeddings.astype("float32"))
 
-# Embedding model pulled directly from the HF model repo.
-encoder = SentenceTransformer(EMB_MODEL_ID)
+# 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.
+encoder = SentenceTransformer(EMB_MODEL_ID, device="cpu")
 print(f"Ready: {index.ntotal:,} stocks | encoder {EMB_MODEL_ID}")
 
-# Generation model is heavy on CPU, so it is loaded on first use only.
-_gen = {}
-
-def _load_generator():
-    if "model" not in _gen:
-        import torch
-        from transformers import AutoTokenizer, AutoModelForCausalLM
-        gid = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
-        print(f"Lazy-loading {gid}...")
-        tok = AutoTokenizer.from_pretrained(gid)
-        tok.padding_side = "left"
-        if tok.pad_token is None:
-            tok.pad_token = tok.eos_token
-        m = AutoModelForCausalLM.from_pretrained(gid, dtype=torch.float32).to("cpu")
-        m.eval()
-        _gen.update(tok=tok, model=m, torch=torch)
-    return _gen["tok"], _gen["model"], _gen["torch"]
+# ---------------------------------------------------------------------------
+# 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.
+# ---------------------------------------------------------------------------
+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
+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}")
 
 # ---------------------------------------------------------------------------
 # QUESTIONNAIRE -> QUERY   (identical logic to the notebook)
@@ -142,7 +150,8 @@ def recommend(risk_level, investment_amount, sector, target_market,
     """
     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.
+    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)
@@ -203,6 +212,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:
@@ -213,7 +223,8 @@ def _tidy(s, max_words=30):
 
 def _stock_prompt(risk, style, r):
     """Qualitative descriptors only — the model never sees a raw figure, so it
-    cannot perform arithmetic on one."""
+    cannot perform arithmetic on one. This is a structural mitigation: given raw
+    numbers, the benchmark 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")
@@ -225,20 +236,30 @@ def _stock_prompt(risk, 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.
+    duration=30 is generous for three short generations; shorter declared
+    durations receive higher queue priority."""
+    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)
+    return gen_tok.batch_decode(out[:, enc["input_ids"].shape[1]:],
+                                skip_special_tokens=True)
+
 def generate_rationale(risk, amount, sector, market, style, recs):
-    tok, model, torch = _load_generator()
-    prompts = [tok.apply_chat_template(
+    """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)}],
         tokenize=False, add_generation_prompt=True) for _, r in recs.iterrows()]
 
-    enc = tok(prompts, return_tensors="pt", padding=True)
-    with torch.no_grad():
-        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)
+    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 "
@@ -260,19 +281,19 @@ 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)
     if recs.empty:
         yield pd.DataFrame(), note, ""
         return
     cols = [c for c in DISPLAY_COLS if c in recs.columns]
-    yield recs[cols], (note + "\n\n⏳ Generating your personalised explanation… "
-                              "(30–90 seconds on free CPU)").strip(), \
+    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}`"
 
 def run_quickstart(name):
-    """Served from the pre-generated cache — instant, no model load."""
+    """Served from the pre-generated cache — instant, uses no GPU quota."""
     entry = QUICKSTART_CACHE[name]
     tbl = pd.DataFrame(entry["table"])
     for c in ["beta", DIV_COL, "cagr_10yr_pct", "price_usd", "similarity",