Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import time
 import torch
+import spaces
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForCausalLM, pipeline
 
 # ── Device ────────────────────────────────────────────────────────────────────
@@ -48,9 +49,23 @@ Final Answer: <answer>
 Be deliberate. Show your full working. Each step should be one clear thought."""
 
 # ── Step Probe inference ──────────────────────────────────────────────────────
-def probe_step(problem: str, steps_so_far: list) -> tuple:
-    context = " ".join(steps_so_far[:-1]) if len(steps_so_far) > 1 else ""
+def is_notation_step(step: str) -> bool:
+    """Detect pure LaTeX/notation steps — restatements, not new reasoning."""
+    s = step.strip()
+    return (
+        s.startswith("\\[") or s.startswith("$$") or
+        s.startswith("\\(") or
+        (s.startswith("\\") and len(s.split()) < 8)
+    )
+
+def probe_step(problem: str, steps_so_far: list, prev_result: tuple = None) -> tuple:
     current = steps_so_far[-1] if steps_so_far else ""
+
+    # Notation-only steps inherit previous score — they restate, not reason
+    if is_notation_step(current):
+        return prev_result if prev_result is not None else ("OK", 0.9)
+
+    context = " ".join(steps_so_far[:-1]) if len(steps_so_far) > 1 else ""
     input_text = (
         f"[PROBLEM] {problem.strip()} "
         f"[STEPS SO FAR] {context} "
@@ -70,6 +85,7 @@ def probe_step(problem: str, steps_so_far: list) -> tuple:
     return label, conf
 
 # ── Core pipeline ─────────────────────────────────────────────────────────────
+@spaces.GPU
 def run_twt(problem: str):
     if not problem.strip():
         yield ("", "<div class='msg warn'>⚠️ Enter a problem to analyze.</div>", "<div class='msg empty'>Waiting...</div>")
@@ -77,7 +93,6 @@ def run_twt(problem: str):
 
     cot_html = ""
 
-    # Build prompt
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user",   "content": problem.strip()}
@@ -95,15 +110,17 @@ def run_twt(problem: str):
 
     lines        = full_response.strip().split("\n")
     parsed_steps = [l.strip() for l in lines if l.strip()]
-    displayed_steps = []
-    first_fail_seen = False
+    displayed_steps  = []
+    first_fail_seen  = False
+    prev_probe_result = None
 
     for i, step in enumerate(parsed_steps):
         displayed_steps.append(step)
         is_final = step.lower().startswith("final answer")
 
         if not is_final:
-            label, conf = probe_step(problem, displayed_steps)
+            label, conf = probe_step(problem, displayed_steps, prev_probe_result)
+            prev_probe_result = (label, conf)
         else:
             label, conf = "FINAL", 1.0
 
@@ -124,18 +141,20 @@ def run_twt(problem: str):
         </div>"""
         cot_html += step_card
 
-        yield (full_response, cot_html, build_trace(problem, displayed_steps, first_fail_seen))
+        yield (full_response, cot_html, build_trace(problem, displayed_steps, first_fail_seen, prev_probe_result))
         time.sleep(0.1)
 
-    yield (full_response, cot_html, build_trace(problem, displayed_steps, first_fail_seen, done=True))
+    yield (full_response, cot_html, build_trace(problem, displayed_steps, first_fail_seen, prev_probe_result, done=True))
 
 
-def build_trace(problem, displayed, first_fail, done=False):
+def build_trace(problem, displayed, first_fail, prev_result=None, done=False):
     faults, rows = 0, ""
+    last_result = None
     for i, s in enumerate(displayed):
         is_final = s.lower().startswith("final answer")
         if not is_final:
-            lbl, conf = probe_step(problem, displayed[:i+1])
+            lbl, conf = probe_step(problem, displayed[:i+1], last_result)
+            last_result = (lbl, conf)
             if lbl == "FAIL": faults += 1
             dot   = f"<span class='dot dot-{'ok' if lbl == 'OK' else 'fail'}'></span>"
             rows += f"<div class='trace-row'>{dot} Step {i+1} <span class='trace-conf'>&mdash; {int(conf*100)}%</span></div>"
@@ -160,34 +179,32 @@ CSS = """
 *, *::before, *::after { box-sizing: border-box; }
 body, .gradio-container, .gradio-container * {
     font-family: 'Inter', -apple-system, sans-serif !important;
-    color: #1e293b !important; /* Forces dark slate text universally */
+    color: #1e293b !important;
 }
 .gradio-container {
-    background-color: #f4f6f8 !important; /* Cool corporate gray backdrop */
-    max-width: 1280px !important; /* Widened slightly to fit dense data */
-    padding-top: 1rem !important; /* Reduced top padding */
+    background-color: #f4f6f8 !important;
+    max-width: 1280px !important;
+    padding-top: 1rem !important;
 }
 
 /* Header */
 .twt-header {
     text-align: center;
-    padding: 0 1rem 1rem; /* Compact padding */
-    margin-bottom: 1.5rem; /* Reduced margin */
+    padding: 0 1rem 1rem;
+    margin-bottom: 1.5rem;
     border-bottom: 1px solid #cbd5e1;
 }
 .twt-title {
-    font-size: 2rem; /* Scaled down for density */
+    font-size: 2rem;
     font-weight: 800;
     letter-spacing: -0.03em;
     font-family: 'JetBrains Mono', monospace !important;
     margin-bottom: 0.25rem;
     color: #0f172a !important;
 }
-.twt-title span {
-    color: #0033a0 !important; /* Morgan Stanley signature blue */
-}
+.twt-title span { color: #0033a0 !important; }
 .twt-sub {
-    font-size: 0.85rem; /* Scaled down */
+    font-size: 0.85rem;
     color: #475569 !important;
     font-weight: 500;
     letter-spacing: 0.02em;
@@ -195,12 +212,12 @@ body, .gradio-container, .gradio-container * {
 
 /* Labels */
 .input-label, .panel-label {
-    font-size: 0.7rem; /* Compact label */
+    font-size: 0.7rem;
     font-weight: 700;
     color: #334155 !important;
     text-transform: uppercase;
     letter-spacing: 0.08em;
-    margin-bottom: 0.5rem; /* Reduced margin */
+    margin-bottom: 0.5rem;
     display: block;
 }
 .panel-label {
@@ -212,11 +229,11 @@ body, .gradio-container, .gradio-container * {
 textarea {
     background: #ffffff !important;
     border: 1px solid #94a3b8 !important;
-    border-radius: 6px !important; /* Sharper corners */
+    border-radius: 6px !important;
     color: #0f172a !important;
-    font-size: 0.85rem !important; /* Scaled down text */
+    font-size: 0.85rem !important;
     line-height: 1.5 !important;
-    padding: 0.75rem !important; /* Compact padding */
+    padding: 0.75rem !important;
     box-shadow: inset 0 1px 2px rgba(15, 23, 42, 0.05) !important;
     transition: all 0.2s ease !important;
 }
@@ -229,13 +246,13 @@ textarea::placeholder { color: #94a3b8 !important; }
 
 /* Primary Button */
 button.primary {
-    background: #0033a0 !important; /* Morgan Stanley Deep Navy */
+    background: #0033a0 !important;
     color: #ffffff !important;
     font-weight: 600 !important;
     border: none !important;
-    border-radius: 6px !important; /* Sharper corners */
-    padding: 0.6rem 1.25rem !important; /* Compact padding */
-    font-size: 0.85rem !important; /* Scaled down */
+    border-radius: 6px !important;
+    padding: 0.6rem 1.25rem !important;
+    font-size: 0.85rem !important;
     box-shadow: 0 2px 4px -1px rgba(0, 51, 160, 0.2) !important;
     transition: all 0.15s ease-in-out !important;
     cursor: pointer !important;
@@ -255,14 +272,14 @@ button.primary:active {
 .step-card {
     display: flex;
     align-items: flex-start;
-    gap: 0.75rem; /* Reduced gap */
-    padding: 0.75rem 1rem; /* Compact padding */
-    margin-bottom: 0.5rem; /* Reduced margin to pack them closer */
+    gap: 0.75rem;
+    padding: 0.75rem 1rem;
+    margin-bottom: 0.5rem;
     background: #ffffff;
     border: 1px solid #cbd5e1;
-    border-radius: 6px; /* Sharper corners */
+    border-radius: 6px;
     line-height: 1.5;
-    animation: fadeUp 0.2s ease-out both; /* Faster animation */
+    animation: fadeUp 0.2s ease-out both;
     box-shadow: 0 1px 2px rgba(15, 23, 42, 0.03);
     transition: all 0.15s ease;
 }
@@ -274,24 +291,22 @@ button.primary:active {
     border-color: #fca5a5;
     background: #fff1f2;
 }
-.step-card.step-fault:hover {
-    border-color: #f87171;
-}
+.step-card.step-fault:hover { border-color: #f87171; }
 @keyframes fadeUp {
     from { opacity: 0; transform: translateY(6px); }
     to   { opacity: 1; transform: translateY(0); }
 }
-.step-icon { font-size: 1rem; margin-top: 1px; } /* Slightly smaller icon */
+.step-icon { font-size: 1rem; margin-top: 1px; }
 .step-text { flex: 1; color: #1e293b !important; font-size: 0.85rem; }
 
 /* Status Badges */
 .badge {
     flex-shrink: 0;
-    font-size: 0.6rem; /* Scaled down */
+    font-size: 0.6rem;
     font-weight: 700;
     font-family: 'JetBrains Mono', monospace !important;
     letter-spacing: 0.02em;
-    padding: 0.15rem 0.4rem; /* Compact padding */
+    padding: 0.15rem 0.4rem;
     border-radius: 4px;
     margin-top: 2px;
     text-transform: uppercase;
@@ -304,37 +319,37 @@ button.primary:active {
 .trace-row {
     display: flex;
     align-items: center;
-    gap: 0.5rem; /* Tighter gap */
-    font-size: 0.75rem; /* Compact font */
+    gap: 0.5rem;
+    font-size: 0.75rem;
     color: #334155 !important;
     font-family: 'JetBrains Mono', monospace !important;
-    padding: 0.35rem 0; /* Reduced padding */
+    padding: 0.35rem 0;
     border-bottom: 1px solid #e2e8f0;
 }
 .trace-conf { color: #64748b !important; }
-.dot { width: 6px; height: 6px; border-radius: 50%; flex-shrink: 0; } /* Smaller dots */
+.dot { width: 6px; height: 6px; border-radius: 50%; flex-shrink: 0; }
 .dot-ok    { background: #10b981; box-shadow: 0 0 0 2px #d1fae5; }
 .dot-fail  { background: #e11d48; box-shadow: 0 0 0 2px #ffe4e6; }
 .dot-final { background: #0033a0; box-shadow: 0 0 0 2px #dbeafe; }
 
 /* Progress Bar */
 .trace-summary {
-    margin-top: 1rem; /* Compact top margin */
+    margin-top: 1rem;
     padding-top: 0.75rem;
     border-top: 1px solid #cbd5e1;
 }
 .trace-label {
-    font-size: 0.65rem; /* Scaled down */
+    font-size: 0.65rem;
     font-weight: 700;
     color: #475569 !important;
     text-transform: uppercase;
     letter-spacing: 0.05em;
-    margin-bottom: 0.4rem; /* Tighter margin */
+    margin-bottom: 0.4rem;
 }
 .trace-bar-bg {
     background: #cbd5e1;
     border-radius: 999px;
-    height: 4px; /* Thinner progress bar */
+    height: 4px;
     overflow: hidden;
 }
 .trace-bar {
@@ -358,7 +373,7 @@ button.primary:active {
     color: #64748b !important;
     font-size: 0.8rem;
     font-family: 'JetBrains Mono', monospace !important;
-    padding: 1.5rem 1rem; /* Compact empty state */
+    padding: 1.5rem 1rem;
     text-align: center;
     background: #ffffff;
     border: 1px dashed #94a3b8;
@@ -375,10 +390,11 @@ HEADER = """
 """
 
 EXAMPLES = [
-    "If a bat and a ball cost $1.10 in total, and the bat costs $1 more than the ball, how much does the ball cost?",
-    "A farmer has 17 sheep. All but 9 die. How many sheep are left?",
-    "What is 15% of 80? Then add that to 25% of 60.",
-    "If you have a 3-gallon jug and a 5-gallon jug, how do you measure exactly 4 gallons?",
+    "Janet has 3 times as many apples as John. John has 8 apples. How many apples does Janet have?",
+    "A store had 50 chocolates. They sold 18 in the morning and 14 in the afternoon. How many chocolates are left?",
+    "Tom reads 12 pages every day. How many pages will he read in 3 weeks?",
+    "A train travels 60 miles per hour. How far will it travel in 2 hours and 30 minutes?",
+    "Lucy earns $15 per hour. She worked 6 hours on Monday and 4 hours on Tuesday. How much did she earn in total?",
 ]
 
 with gr.Blocks(css=CSS, title="ThinkWhileThinking") as demo:
@@ -401,4 +417,4 @@ with gr.Blocks(css=CSS, title="ThinkWhileThinking") as demo:
     run_btn.click(fn=run_twt, inputs=[problem_input], outputs=[raw_output, cot_output, trace_output])
     problem_input.submit(fn=run_twt, inputs=[problem_input], outputs=[raw_output, cot_output, trace_output])
 
-demo.launch()
+demo.launch()