Update app.py

app.py

@@ -41,7 +41,8 @@ import seaborn as sns  # noqa: E402
 from matplotlib.ticker import (FuncFormatter, LogLocator,  # noqa: E402
                                NullFormatter, ScalarFormatter)
 from huggingface_hub import HfApi, hf_hub_download  # noqa: E402
-from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: E402
+from transformers import (AutoModelForCausalLM, AutoTokenizer,  # noqa: E402
+                          StoppingCriteria, StoppingCriteriaList)
 
 # --------------------------------------------------------------------------- #
 # Model — must be placed on cuda at module level for ZeroGPU
@@ -64,19 +65,57 @@ print(f"[startup] device={DEVICE} dtype={DTYPE}")
 PLOTS_DIR = "plots"
 
 
+class StopOnClosingFence(StoppingCriteria):
+    """Stop once a fenced block has been opened and closed.
+
+    Decoding the tail is more reliable than matching token ids: ``` tokenizes
+    differently depending on what precedes it. Decoding on every step would be
+    wasteful, so we only check every `check_every` tokens.
+    """
+
+    def __init__(self, tokenizer, prompt_len, fences=2, check_every=8):
+        self.tok, self.prompt_len = tokenizer, prompt_len
+        self.fences, self.check_every = fences, check_every
+        self._step = 0
+
+    def __call__(self, input_ids, scores, **kwargs):
+        self._step += 1
+        done = False
+        if self._step % self.check_every == 0:
+            text = self.tok.decode(input_ids[0][self.prompt_len:],
+                                   skip_special_tokens=True)
+            done = text.count("```") >= self.fences
+        # one flag per batch row: StoppingCriteriaList ORs the results together
+        return torch.full((input_ids.shape[0],), done,
+                          dtype=torch.bool, device=input_ids.device)
+
+
 @torch.inference_mode()
-def generate(system: str, user: str, max_new_tokens: int = 900) -> str:
-    """Single entry point to the LLM. Returns only the newly generated text."""
+def generate(system: str, user: str, max_new_tokens: int = 900,
+             stop_fences: int = 0) -> str:
+    """Single entry point to the LLM. Returns only the newly generated text.
+
+    stop_fences=2 halts as soon as the closing ``` of a code block is emitted,
+    instead of running on to max_new_tokens with trailing prose.
+    """
     msgs = [{"role": "system", "content": system}, {"role": "user", "content": user}]
     text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
     inputs = tok(text, return_tensors="pt").to(model.device)
+    prompt_len = inputs["input_ids"].shape[1]
+
+    criteria = None
+    if stop_fences:
+        criteria = StoppingCriteriaList(
+            [StopOnClosingFence(tok, prompt_len, fences=stop_fences)])
+
     out = model.generate(
         **inputs,
         max_new_tokens=max_new_tokens,
         do_sample=False,                      # greedy -> reproducible
         pad_token_id=tok.pad_token_id,
+        stopping_criteria=criteria,
     )
-    return tok.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+    return tok.decode(out[0][prompt_len:], skip_special_tokens=True)
 
 
 # --------------------------------------------------------------------------- #
@@ -354,7 +393,8 @@ def format_error(err: str, code: str) -> str:
 # --------------------------------------------------------------------------- #
 # 5. Deterministic facts and tables
 # --------------------------------------------------------------------------- #
-REDUNDANT_R = 0.99      # |r| at or above this is a duplicate encoding, not a finding
+REDUNDANT_R = 0.99        # |r| at or above this is a duplicate encoding, not a finding
+MAX_GROUP_BLOCKS = 10     # one-hot-heavy datasets otherwise emit dozens of tables
 
 
 def fmt_num(v):
@@ -415,7 +455,13 @@ def make_tables(df: pd.DataFrame, max_card: int = 6) -> dict:
         for c in df.columns if df[c].nunique(dropna=True) <= max_card
         for b in bins if b != c
     ]
-    t["groups"] = "\n\n".join(blocks) if blocks else "_No low-cardinality groupings._"
+    if blocks:
+        t["groups"] = "\n\n".join(blocks[:MAX_GROUP_BLOCKS])
+        if len(blocks) > MAX_GROUP_BLOCKS:
+            t["groups"] += (f"\n\n_+{len(blocks) - MAX_GROUP_BLOCKS} more group "
+                            "comparisons omitted._")
+    else:
+        t["groups"] = "_No low-cardinality groupings._"
 
     real, dup = corr_pairs(df)
     t["corr"] = (pd.DataFrame(real[:8]).to_markdown(index=False)
@@ -632,7 +678,8 @@ def _pipeline(instruction: str, dataset: str):
            *BLANK)
 
     code = sanitize(extract_code(
-        generate(CODE_SYSTEM, f"{context}\n\nTask: {instruction}", max_new_tokens=1200)))
+        generate(CODE_SYSTEM, f"{context}\n\nTask: {instruction}",
+                 max_new_tokens=1200, stop_fences=2)))
     res = run_code(code, df)
     yield (f"Attempt 1: {'ok' if res['ok'] else 'failed'}. Executing …",
            None, code, res["stdout"], res["plots"], "", "")
@@ -643,7 +690,8 @@ def _pipeline(instruction: str, dataset: str):
                res["plots"], "", "")
         fix = (f"{context}\n\nThis code failed:\n```python\n{code}\n```\n\n"
                f"Error:\n{format_error(res['error'], code)}\n\nReturn the corrected script.")
-        code = sanitize(extract_code(generate(FIX_SYSTEM, fix, max_new_tokens=1200)))
+        code = sanitize(extract_code(
+            generate(FIX_SYSTEM, fix, max_new_tokens=1200, stop_fences=2)))
         res = run_code(code, df)
 
     plots = sorted(glob.glob(f"{PLOTS_DIR}/*.png"))
@@ -698,6 +746,7 @@ def run_agent(instruction: str, dataset: str):
 # --------------------------------------------------------------------------- #
 DEFAULT_INSTRUCTION = ("Run a comprehensive exploratory data analysis, highlighting "
                        "missing values, distributions, and key feature correlations.")
+DEFAULT_DATASET = "Kogann/stockmatch-synthetic"
 
 QS_DIR = "quickstarts"
 QUICKSTARTS = [
@@ -706,10 +755,10 @@ QUICKSTARTS = [
      "blurb": "77% of `cabin` is missing, but the raw file hides it as `''`.",
      "dataset": "mstz/titanic",
      "instruction": DEFAULT_INSTRUCTION},
-    {"slug": "iris",
-     "label": "🌸 Iris — a clean baseline",
-     "blurb": "No missing values, a row-id column the agent drops, tight correlations.",
-     "dataset": "scikit-learn/iris",
+    {"slug": "stockmatch",
+     "label": "📈 StockMatch — synthetic stocks",
+     "blurb": "12,500 generated stocks — does synthetic data carry real structure?",
+     "dataset": "Kogann/stockmatch-synthetic",
      "instruction": "Explore this dataset: missing values, distributions, correlations."},
     {"slug": "housing",
      "label": "🏠 Canada housing — skewed prices",
@@ -756,7 +805,7 @@ def quickstart_handler(spec: dict):
 
 
 def reset_form():
-    return DEFAULT_INSTRUCTION, "mstz/titanic"
+    return DEFAULT_INSTRUCTION, DEFAULT_DATASET
 
 
 # --------------------------------------------------------------------------- #
@@ -797,7 +846,7 @@ with gr.Blocks(title="EDAgent", theme=gr.themes.Soft(), css=CSS) as demo:
     with gr.Row():
         instruction = gr.Textbox(label="Prompt instruction", value=DEFAULT_INSTRUCTION,
                                  lines=3, scale=3)
-        dataset = gr.Textbox(label="Hugging Face dataset id", value="mstz/titanic",
+        dataset = gr.Textbox(label="Hugging Face dataset id", value=DEFAULT_DATASET,
                              placeholder="owner/name", lines=1, scale=1)
     with gr.Row():
         run_btn = gr.Button("🚀 Run EDA Agent", variant="primary", size="lg", scale=4)