Update app.py

app.py

@@ -253,15 +253,38 @@ def _stock_prompt(risks, style, r):
 
 @spaces.GPU(duration=30)
 def _generate_sentences(prompts):
-    """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,
-                                 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)
+    """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()}
+
+
+@spaces.GPU(duration=10)
+def _gpu_selftest():
+    """Diagnostic: does ZeroGPU work at all, independent of our model?
+    Isolates an infrastructure problem from a model-loading problem."""
+    import traceback
+    try:
+        x = torch.randn(4, 4).to("cuda")
+        return f"✅ GPU OK — {torch.cuda.get_device_name(0)}, tensor on {x.device}"
+    except Exception:
+        return "❌ GPU FAILED\n```\n" + traceback.format_exc()[-1200:] + "\n```"
 
 def generate_rationale(risks, amount, sectors, markets, style, recs):
     """One sentence per stock, generated in a single batched pass. All factual
@@ -272,11 +295,12 @@ def generate_rationale(risks, amount, sectors, markets, style, recs):
          {"role": "user",   "content": _stock_prompt(risks, style, r)}],
         tokenize=False, add_generation_prompt=True) for _, r in recs.iterrows()]
 
-    try:
-        sents = _generate_sentences(prompts)
-    except Exception as e:
-        return (f"_Recommendations above are complete. The written explanation "
-                f"could not be generated._\n\n`{type(e).__name__}: {e}`")
+    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", "no traceback")[-1500:] + "\n```")
+    sents = result["sentences"]
 
     amt = parse_amount(amount)
     rl  = "/".join(_as_list(risks)).lower() or "flexible"
@@ -287,7 +311,7 @@ def generate_rationale(risks, amount, sectors, markets, style, recs):
     lines = [f"Based on your {rl}-risk profile and {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(amt // r["price_usd"]) if (amt > 0 and r["price_usd"]) else None
+        sh  = int(amt // r["price_usd"]) if (amt > 0 and r["price_usd"]) else None
         aff = f" ({sh} shares affordable)" if sh else ""
         lines.append(f"{i}. **{r['ticker']}** — beta {r['beta']:.2f}, "
                      f"{r[DIV_COL]:.2f}% yield, {r['cagr_10yr_pct']:.1f}% growth, "
@@ -389,6 +413,10 @@ 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()
@@ -397,5 +425,6 @@ 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()