Update app.py

app.py

@@ -8,12 +8,11 @@ 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 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.
+The three Quick Starters are served from a pre-generated cache and never touch
+the language model, so the common path is instant.
 """
 
-import os, json, re
+import os, json, re, traceback
 import numpy as np
 import pandas as pd
 import gradio as gr
@@ -58,22 +57,28 @@ index.add(embeddings.astype("float32"))
 encoder = SentenceTransformer(EMB_MODEL_ID, device="cpu")
 print(f"Ready: {index.ntotal:,} stocks | encoder {EMB_MODEL_ID}")
 
-# 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.
+# ---------------------------------------------------------------------------
+# GENERATION MODEL
+# Loaded at STARTUP rather than on first request. Lazy loading pushed a ~60s
+# download-and-load into the first user's request, which exceeded the request
+# timeout and surfaced as a bare "Error". Paying the cost once at boot keeps
+# every request fast. If it fails, the app still serves recommendations —
+# retrieval does not depend on the language model.
+# ---------------------------------------------------------------------------
 GEN_MODEL_ID = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
-_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"]
+gen_tok, gen_model, GEN_LOAD_ERROR = None, None, None
+try:
+    print(f"Loading {GEN_MODEL_ID} ...")
+    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.float32)
+    gen_model.eval()
+    print("Generation model ready.")
+except Exception as e:
+    GEN_LOAD_ERROR = f"{type(e).__name__}: {e}"
+    print("Generation model FAILED to load:", GEN_LOAD_ERROR)
 
 # ---------------------------------------------------------------------------
 # QUESTIONNAIRE -> QUERY
@@ -147,8 +152,6 @@ def build_query(risks, sectors, markets, style, caps):
     if risk_phrase:
         parts.append(risk_phrase)
     q = ". ".join([p for p in parts if p])
-    # If nothing at all was selected, fall back to the project's core intent so
-    # the search still has something meaningful to match against.
     return q or "a stock that helps protect purchasing power against inflation"
 
 def embed_query(text):
@@ -264,19 +267,17 @@ def _stock_prompt(risks, style, r):
             f"investor and note one caveat.")
 
 def _generate_sentences(prompts):
-    """Runs on CPU. Errors are returned as data rather than raised, so a
-    generation failure never takes down the recommendation table."""
-    import traceback
+    """Errors are returned as data rather than raised, so a generation failure
+    never takes down the recommendation table."""
     try:
-        tok, model = _load_generator()
-        enc = tok(prompts, return_tensors="pt", padding=True)
+        enc = gen_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)
+            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()}
@@ -285,18 +286,24 @@ 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."""
-    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()]
-
-    result = _generate_sentences(prompts)
-    if not result.get("ok"):
+    if gen_model is None:
+        return ("_Recommendations above are complete. The language model is "
+                f"unavailable._\n\n`{GEN_LOAD_ERROR}`")
+    try:
+        prompts = [gen_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()]
+        result = _generate_sentences(prompts)
+        if not result.get("ok"):
+            return ("_Recommendations above are complete. The written explanation "
+                    "could not be generated._\n\n```\n"
+                    + result.get("traceback", "")[-1200:] + "\n```")
+        sents = result["sentences"]
+    except Exception:
         return ("_Recommendations above are complete. The written explanation "
                 "could not be generated._\n\n```\n"
-                + result.get("traceback", "no traceback")[-1500:] + "\n```")
-    sents = result["sentences"]
+                + traceback.format_exc()[-1200:] + "\n```")
 
     amt = parse_amount(amount)
     rl  = "/".join(_as_list(risks)).lower() or "flexible"
@@ -353,19 +360,18 @@ def _prettify(frame):
 # ---------------------------------------------------------------------------
 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."""
+    generation of some seconds does not look like a frozen page."""
     recs, note, _query = recommend(risks, amount, sectors, markets, style, caps)
     if recs.empty:
         yield pd.DataFrame(), note
         return
     pretty = _prettify(recs)
-    yield pretty, (note + "\n⏳ Writing your personalised explanation… "
-                          "(about a minute on free CPU)").strip()
+    yield pretty, (note + "\n⏳ Writing your personalised explanation…").strip()
     text = generate_rationale(risks, amount, sectors, markets, style, recs)
     yield pretty, (note + "\n" + text).strip()
 
 def run_quickstart(name):
-    """Served from the pre-generated cache — instant, no model load."""
+    """Served from the pre-generated cache — instant, no model call."""
     entry = QUICKSTART_CACHE[name]
     tbl = pd.DataFrame(entry["table"])
     for c in ["beta", DIV_COL, "cagr_10yr_pct", "price_usd", "similarity",