app.py

"""
StockMatch AI — Gradio application (Part 5)

QUESTIONNAIRE -> EMBEDDING -> FAISS RETRIEVAL -> FILTERS -> GENERATED RATIONALE

Data sources, per the assignment constraints:
  * 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 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, traceback
import numpy as np
import pandas as pd
import gradio as gr
import faiss
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"

# ---------------------------------------------------------------------------
# LOAD ARTIFACTS
# ---------------------------------------------------------------------------
print("Loading artifacts...")
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"

# 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}")
else:
    embeddings = np.load(hf_hub_download(DATASET_REPO, EMB_FILE, repo_type="dataset"))
    print(f"Embeddings loaded from dataset repo: {EMB_FILE}")

df = pd.read_parquet(hf_hub_download(DATASET_REPO, "stock_metadata.parquet",
                                     repo_type="dataset"))
QUICKSTART_CACHE = json.load(open(hf_hub_download(
    DATASET_REPO, "quickstart_rationales.json", repo_type="dataset")))

assert len(embeddings) == len(df), "embeddings and metadata are misaligned"

# Vectors are L2-normalised, so INNER PRODUCT is identical to cosine similarity.
index = faiss.IndexFlatIP(embeddings.shape[1])
index.add(embeddings.astype("float32"))

encoder = SentenceTransformer(EMB_MODEL_ID, device="cpu")
print(f"Ready: {index.ntotal:,} stocks | encoder {EMB_MODEL_ID}")

# ---------------------------------------------------------------------------
# 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_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
    # transformers 4.44 expects torch_dtype=, not dtype=. float32 is the CPU
    # default anyway, so the argument is simply omitted.
    gen_model = AutoModelForCausalLM.from_pretrained(GEN_MODEL_ID)
    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
# Questions 1, 3, 4 and 6 accept MULTIPLE selections, and all six start EMPTY:
# an unanswered question means "no restriction" rather than a hidden default,
# so the app never silently filters on something the user did not choose.
# ---------------------------------------------------------------------------
RISK_TEXT = {
    "Low":    "low volatility, defensive, conservative risk, capital preservation",
    "Medium": "moderate volatility, balanced risk",
    "High":   "high volatility, aggressive risk",
}
STYLE_TEXT = {
    "Dividend": "high dividend income, strong payout, steady historical growth",
    "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"}

# 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",
    "Japan": "Asia", "China": "Asia", "Asia": "Asia",
}
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):
    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 parse_amount(a):
    """The dropdown shows formatted currency; strip it back to a number.
    Returns 0 when unanswered, which downstream treats as 'no budget limit'."""
    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):
    """Risk phrases are repeated at the end: 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))
    if risk_phrase:
        parts.append(risk_phrase)
    q = ". ".join([p for p in parts if p])
    return q or "a stock that helps protect purchasing power against inflation"

def embed_query(text):
    return encoder.encode([QUERY_PREFIX + text],
                          normalize_embeddings=True).astype("float32")

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(risks, sectors, markets, style, caps)
    regions = parse_markets(markets)
    sec     = _as_list(sectors)
    amt     = parse_amount(amount)

    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[_beta_mask(res, risks)]
    if amt > 0:                                  # unanswered budget = no limit
        res = res[res["price_usd"] <= amt]
    if sec:
        res = res[res["sector"].isin(sec)]
    if regions:
        res = res[res["market_region"].isin(regions)]
    if style == "Dividend":
        res = res[res[DIV_COL] > 0]

    note = ""
    if res.empty:
        # 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 sec:
            pool = pool[pool["sector"].isin(sec)]
        if regions:
            pool = pool[pool["market_region"].isin(regions)]
        if amt > 0:
            pool = pool[pool["price_usd"] <= amt]
        if style == "Dividend":
            pool = pool[pool[DIV_COL] > 0]
        if not pool.empty:
            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 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.\n")
            res = pool
        else:
            note = ("⚠️ No stock matches these constraints. "
                    "Try widening your budget, sectors or markets.")
            res = df.head(0)
    return res.head(top_k).reset_index(drop=True), note, query

# ---------------------------------------------------------------------------
# GENERATION   (facts from code, meaning from the model)
# ---------------------------------------------------------------------------
STOCK_SYSTEM_PROMPT = (
    "You explain what a stock's figures MEAN for one investor. "
    "Write ONE short sentence, maximum 20 words. "
    "Do NOT repeat the numbers — they are already displayed to the user. "
    "Do NOT calculate anything. Do NOT invent a company name. "
    "Do NOT mention any index, benchmark or other company."
)
_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:
        s = " ".join(w[:max_words])
    if s.split() and s.split()[-1].lower().strip(",") in _DANGLING and "," in s:
        s = s.rsplit(",", 1)[0]
    return s.rstrip(",;: ") + "."

def _stock_prompt(risks, style, r):
    """Qualitative descriptors only — the model never sees a raw figure, so it
    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")
    rl = "/".join(_as_list(risks)).lower() or "flexible"
    st = (style or "balanced").lower()
    return (f"Investor: {rl} risk tolerance, prefers {st} 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 "
            f"investor and note one caveat.")

def _generate_sentences(prompts):
    """Errors are returned as data rather than raised, so a generation failure
    never takes down the recommendation table."""
    try:
        enc = gen_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)
        return {"ok": True, "sentences": sents}
    except Exception:
        return {"ok": False, "traceback": traceback.format_exc()}

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."""
    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"
                + traceback.format_exc()[-1200:] + "\n```")

    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"
    budget = f"${amt:,.0f} budget" if amt > 0 else "no set budget"

    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
        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, "
                     f"${r['price_usd']:,.2f} per share{aff}. {_tidy(s)}")
    lines += ["", "_This dataset is synthetic and was generated for an educational "
                  "project. It is not investment advice._"]
    return "\n".join(lines)

# ---------------------------------------------------------------------------
# RESULTS TABLE FORMATTING
# 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"]

COLUMN_LABELS = {
    "ticker":        "Ticker",
    "sector":        "Sector",
    "market_region": "Region",
    "cluster_name":  "Segment",
    "risk_level":    "Risk",
    "beta":          "Volatility",
    DIV_COL:         "Dividend %",
    "cagr_10yr_pct": "Growth %",
    "price_usd":     "Price $",
    "similarity":    "Match %",
}

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"].astype("float64") * 100).round(1)
    for c in ["beta", DIV_COL, "cagr_10yr_pct", "price_usd"]:
        if c in out:
            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
    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…").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 call."""
    entry = QUICKSTART_CACHE[name]
    tbl = pd.DataFrame(entry["table"])
    for c in ["beta", DIV_COL, "cagr_10yr_pct", "price_usd", "similarity",
              "market_cap_millions_usd"]:
        if c in tbl.columns:
            tbl[c] = pd.to_numeric(tbl[c], errors="coerce")
    return _prettify(tbl), entry["rationale"]

def reset_form():
    """Return every question to its unanswered state. The order must match the
    output list on the reset click handler."""
    return [], None, [], [], None, []

# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="StockMatch AI", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# 📈 StockMatch AI\n"
        "**Money sitting in a savings account loses value every year. "
        "Find stocks matched to your risk tolerance, budget and goals.**\n\n"
        "Answer six questions, or click a preset profile. Matches come from a FAISS "
        "search over 12,500 synthetic stocks embedded with `BAAI/bge-small-en-v1.5`; "
        "the written explanation is generated by `SmolLM2-1.7B-Instruct`.\n\n"
        "⚠️ *All data is synthetic and generated for an educational project. "
        "This is not investment advice.*"
    )

    gr.Markdown("### ⚡ Quick Starters — one-click example investors (instant)")
    with gr.Row():
        b1 = gr.Button("👵 Cautious Retiree", variant="secondary")
        b2 = gr.Button("👔 Balanced Professional", variant="secondary")
        b3 = gr.Button("🚀 Young Growth Seeker", variant="secondary")

    gr.Markdown("### 📋 Or build your own profile\n"
                "Questions 1, 3, 4 and 6 accept multiple answers. "
                "Leave any question blank to apply no restriction.")
    with gr.Row():
        with gr.Column():
            q1 = gr.CheckboxGroup(["Low", "Medium", "High"], value=[],
                                  label="1. Risk level")
            q2 = gr.Dropdown(AMOUNTS, value=None,
                             label="2. Investment amount (USD)")
            q3 = gr.Dropdown(SECTORS, value=[], multiselect=True,
                             label="3. Sector")
        with gr.Column():
            q4 = gr.Dropdown(MARKETS, value=[], multiselect=True,
                             label="4. Target market (grouped into US / Europe / "
                                   "Israel / Asia)")
            q5 = gr.Radio(["Dividend", "Growth"], value=None,
                          label="5. Stock type preference")
            q6 = gr.CheckboxGroup(["Small", "Medium", "Large"], value=[],
                                  label="6. Company size")

    with gr.Row():
        go    = gr.Button("🔍 Find my stocks", variant="primary", size="lg", scale=4)
        reset = gr.Button("↺ Reset", variant="secondary", size="lg", scale=1)

    gr.Markdown("### Results")
    table = gr.Dataframe(label="Your matches", interactive=False, wrap=True)
    text  = gr.Markdown()

    go.click(run_custom, [q1, q2, q3, q4, q5, q6], [table, text])
    reset.click(reset_form, None, [q1, q2, q3, q4, q5, q6])
    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])

demo.launch()