Update app.py

app.py

@@ -4,7 +4,7 @@ import random
 from huggingface_hub import InferenceClient
 
 # ── Model config ──────────────────────────────────────────────────────────────
-REASONER_MODEL = "Qwen/Qwen2.5-7B-Instruct"          # swap later for local
+REASONER_MODEL = "Qwen/Qwen2.5-7B-Instruct"
 client = InferenceClient(REASONER_MODEL)
 
 SYSTEM_PROMPT = """You are a careful step-by-step reasoner. 
@@ -17,24 +17,20 @@ Final Answer: <answer>
 
 Be deliberate. Show your full working. Each step should be one clear thought."""
 
-# ── Sentinel (placeholder — will be replaced with fine-tuned model) ───────────
-def sentinel_score(problem: str, steps_so_far: list[str]) -> tuple[str, float]:
+# ── Step Probe (placeholder — will be replaced with fine-tuned model) ─────────
+def probe_step(problem: str, steps_so_far: list[str]) -> tuple[str, float]:
     """
-    Placeholder sentinel. Returns (label, confidence).
+    Step-level reasoning probe. Returns (label, confidence).
     Real version: fine-tuned classifier on PRM800K step-level annotations.
     """
     context = " ".join(steps_so_far).lower()
-    red_flags = [
-        "therefore", "obviously", "clearly", "must be", "so",
-        "=", "thus", "hence"
-    ]
-    # Heuristic: suspiciously fast conclusions or leaps get flagged
+    red_flags = ["therefore", "obviously", "clearly", "must be", "so", "=", "thus", "hence"]
     leap_score = sum(1 for f in red_flags if f in context) / max(len(steps_so_far), 1)
     confidence = round(min(0.95, max(0.4, 0.85 - leap_score * 0.15 + random.uniform(-0.05, 0.05))), 2)
     label = "OK" if confidence >= 0.6 else "FAIL"
     return label, confidence
 
-# ── Core reasoning pipeline ───────────────────────────────────────────────────
+# ── Core pipeline ─────────────────────────────────────────────────────────────
 def run_twt(problem: str):
     if not problem.strip():
         yield (
@@ -44,11 +40,8 @@ def run_twt(problem: str):
         )
         return
 
-    steps_so_far = []
     cot_html = ""
-    trace_html = ""
 
-    # Stream CoT from reasoner
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user", "content": problem.strip()}
@@ -72,15 +65,9 @@ def run_twt(problem: str):
         )
         return
 
-    # Parse into steps
     lines = full_response.strip().split("\n")
-    parsed_steps = []
-    for line in lines:
-        line = line.strip()
-        if line:
-            parsed_steps.append(line)
+    parsed_steps = [l.strip() for l in lines if l.strip()]
 
-    # Animate step by step with sentinel scoring
     displayed_steps = []
     first_fail_seen = False
 
@@ -89,11 +76,10 @@ def run_twt(problem: str):
         is_final = step.lower().startswith("final answer")
 
         if not is_final:
-            label, conf = sentinel_score(problem, displayed_steps)
+            label, conf = probe_step(problem, displayed_steps)
         else:
             label, conf = "FINAL", 1.0
 
-        # Build step card
         if is_final:
             icon = "🏁"
             badge_class = "badge-final"
@@ -110,42 +96,36 @@ def run_twt(problem: str):
                 first_fail_seen = True
 
         step_card = f"""
-        <div class='step-card {"step-fault" if label == "FAIL" else ""}' style='animation-delay:{i*0.05}s'>
+        <div class='step-card {"step-fault" if label == "FAIL" else ""}'>
             <span class='step-icon'>{icon}</span>
             <span class='step-text'>{step}</span>
             <span class='badge {badge_class}'>{badge_text}</span>
         </div>"""
         cot_html += step_card
 
-        # Build trace panel
-        trace_html = build_trace(parsed_steps[:i+1], displayed_steps, first_fail_seen)
-
+        trace_html = build_trace(displayed_steps, first_fail_seen)
         yield (full_response, cot_html, trace_html)
         time.sleep(0.15)
 
-    # Final trace summary
-    fault_count = sum(1 for s in displayed_steps
-                      if sentinel_score(problem, displayed_steps[:displayed_steps.index(s)+1])[0] == "FAIL")
-    yield (full_response, cot_html, build_trace(parsed_steps, displayed_steps, first_fail_seen, done=True))
+    yield (full_response, cot_html, build_trace(displayed_steps, first_fail_seen, done=True))
 
 
-def build_trace(all_steps, displayed, first_fail, done=False):
-    total = len(all_steps)
+def build_trace(displayed, first_fail, done=False):
     faults = 0
     rows = ""
     for i, s in enumerate(displayed):
         is_final = s.lower().startswith("final answer")
         if not is_final:
-            lbl, conf = sentinel_score("", displayed[:i+1])
+            lbl, conf = probe_step("", displayed[:i+1])
             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} — {int(conf*100)}% confidence</div>"
+            rows += f"<div class='trace-row'>{dot} Step {i+1} &mdash; {int(conf*100)}%</div>"
         else:
-            rows += f"<div class='trace-row'><span class='dot dot-final'></span> {s}</div>"
+            rows += f"<div class='trace-row'><span class='dot dot-final'></span> Final Answer</div>"
 
     health = max(0, 100 - faults * 20)
-    bar_color = "#4ade80" if health > 60 else "#f87171"
+    bar_color = "#16a34a" if health > 60 else "#dc2626"
     summary = f"""
     <div class='trace-summary'>
         <div class='trace-label'>Reasoning Health</div>
@@ -159,171 +139,215 @@ def build_trace(all_steps, displayed, first_fail, done=False):
     return f"{rows}{summary}"
 
 
-# ── CSS ───────────────────────────────────────────────────────────────────────
+# ── CSS — clean light theme, high contrast ────────────────────────────────────
 CSS = """
-@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600&family=JetBrains+Mono:wght@400;500&display=swap');
+@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap');
 
-* { box-sizing: border-box; margin: 0; padding: 0; }
+*, *::before, *::after { box-sizing: border-box; }
 
-body, .gradio-container {
-    background: #1e1e2e !important;
-    color: #cdd6f4 !important;
+body, .gradio-container, .gradio-container * {
     font-family: 'Inter', sans-serif !important;
 }
 
-/* Header */
+.gradio-container {
+    background: #f8fafc !important;
+    max-width: 1200px !important;
+}
+
+/* ── Header ── */
 .twt-header {
     text-align: center;
-    padding: 2.5rem 1rem 1.5rem;
-    border-bottom: 1px solid #313244;
+    padding: 2rem 1rem 1.25rem;
+    border-bottom: 2px solid #e2e8f0;
     margin-bottom: 1.5rem;
+    background: #ffffff;
+    border-radius: 12px 12px 0 0;
 }
 .twt-title {
     font-size: 2rem;
-    font-weight: 600;
-    letter-spacing: -0.03em;
-    color: #cdd6f4;
-    font-family: 'JetBrains Mono', monospace;
+    font-weight: 700;
+    letter-spacing: -0.04em;
+    color: #0f172a;
+    font-family: 'JetBrains Mono', monospace !important;
 }
-.twt-title span { color: #cba6f7; }
+.twt-title span { color: #7c3aed; }
 .twt-sub {
-    font-size: 0.8rem;
-    color: #6c7086;
-    margin-top: 0.4rem;
-    letter-spacing: 0.08em;
+    font-size: 0.75rem;
+    color: #94a3b8;
+    margin-top: 0.35rem;
+    letter-spacing: 0.1em;
     text-transform: uppercase;
+    font-weight: 500;
 }
 
-/* Input area */
+/* ── Input ── */
 .input-label {
-    font-size: 0.75rem;
-    font-weight: 500;
-    color: #a6adc8;
-    letter-spacing: 0.06em;
+    font-size: 0.7rem;
+    font-weight: 700;
+    color: #64748b;
+    letter-spacing: 0.1em;
     text-transform: uppercase;
-    margin-bottom: 0.5rem;
+    margin-bottom: 0.4rem;
 }
+
 textarea {
-    background: #181825 !important;
-    border: 1px solid #45475a !important;
+    background: #ffffff !important;
+    border: 1.5px solid #cbd5e1 !important;
     border-radius: 8px !important;
-    color: #cdd6f4 !important;
-    font-family: 'Inter', sans-serif !important;
+    color: #0f172a !important;
     font-size: 0.95rem !important;
-    transition: border-color 0.2s !important;
+    line-height: 1.6 !important;
+    transition: border-color 0.15s !important;
+    padding: 0.75rem !important;
 }
 textarea:focus {
-    border-color: #cba6f7 !important;
+    border-color: #7c3aed !important;
+    box-shadow: 0 0 0 3px rgba(124,58,237,0.1) !important;
     outline: none !important;
 }
+textarea::placeholder { color: #94a3b8 !important; }
 
-/* Button */
+/* ── Button ── */
 button.primary {
-    background: #cba6f7 !important;
-    color: #1e1e2e !important;
+    background: #7c3aed !important;
+    color: #ffffff !important;
     font-weight: 600 !important;
     font-family: 'Inter', sans-serif !important;
     border: none !important;
     border-radius: 8px !important;
     padding: 0.65rem 1.5rem !important;
     font-size: 0.9rem !important;
-    letter-spacing: 0.02em !important;
-    transition: background 0.2s !important;
+    letter-spacing: 0.01em !important;
+    transition: background 0.15s, transform 0.1s !important;
     cursor: pointer !important;
+    width: 100% !important;
+}
+button.primary:hover {
+    background: #6d28d9 !important;
+    transform: translateY(-1px) !important;
+}
+button.primary:active { transform: translateY(0) !important; }
+
+/* ── Panel labels ── */
+.panel-label {
+    font-size: 0.68rem;
+    font-weight: 700;
+    color: #64748b;
+    text-transform: uppercase;
+    letter-spacing: 0.12em;
+    margin-bottom: 0.75rem;
+    padding-bottom: 0.5rem;
+    border-bottom: 2px solid #e2e8f0;
 }
-button.primary:hover { background: #b4befe !important; }
 
-/* Step cards */
+/* ── Step cards ── */
 .step-card {
     display: flex;
     align-items: flex-start;
     gap: 0.75rem;
     padding: 0.75rem 1rem;
     margin-bottom: 0.5rem;
-    background: #181825;
-    border: 1px solid #313244;
+    background: #ffffff;
+    border: 1.5px solid #e2e8f0;
     border-radius: 8px;
-    font-size: 0.88rem;
-    line-height: 1.6;
-    animation: fadeSlide 0.25s ease both;
+    font-size: 0.875rem;
+    line-height: 1.65;
+    animation: fadeUp 0.2s ease both;
+    box-shadow: 0 1px 3px rgba(0,0,0,0.04);
 }
 .step-card.step-fault {
-    border-color: #f38ba8;
-    background: #2a1520;
+    border-color: #fca5a5;
+    background: #fff5f5;
 }
-@keyframes fadeSlide {
-    from { opacity: 0; transform: translateY(6px); }
+@keyframes fadeUp {
+    from { opacity: 0; transform: translateY(8px); }
     to   { opacity: 1; transform: translateY(0); }
 }
-.step-icon { font-size: 1rem; flex-shrink: 0; margin-top: 1px; }
-.step-text { flex: 1; color: #cdd6f4; font-family: 'Inter', sans-serif; }
+.step-icon { font-size: 1rem; flex-shrink: 0; margin-top: 2px; }
+.step-text {
+    flex: 1;
+    color: #1e293b;
+    font-family: 'Inter', sans-serif !important;
+    font-size: 0.875rem;
+}
 .badge {
     flex-shrink: 0;
-    font-size: 0.65rem;
-    font-weight: 600;
-    font-family: 'JetBrains Mono', monospace;
-    letter-spacing: 0.05em;
-    padding: 0.2rem 0.5rem;
+    font-size: 0.62rem;
+    font-weight: 700;
+    font-family: 'JetBrains Mono', monospace !important;
+    letter-spacing: 0.04em;
+    padding: 0.18rem 0.5rem;
     border-radius: 4px;
     align-self: flex-start;
-    margin-top: 2px;
+    margin-top: 3px;
 }
-.badge-ok   { background: #1e3a2a; color: #a6e3a1; border: 1px solid #a6e3a1; }
-.badge-fail { background: #3a1e22; color: #f38ba8; border: 1px solid #f38ba8; }
-.badge-final { background: #2a1e3a; color: #cba6f7; border: 1px solid #cba6f7; }
+.badge-ok   { background: #dcfce7; color: #15803d; border: 1px solid #86efac; }
+.badge-fail { background: #fee2e2; color: #b91c1c; border: 1px solid #fca5a5; }
+.badge-final { background: #ede9fe; color: #6d28d9; border: 1px solid #c4b5fd; }
 
-/* Trace panel */
+/* ── Trace panel ── */
 .trace-row {
     display: flex;
     align-items: center;
     gap: 0.6rem;
-    font-size: 0.8rem;
-    color: #a6adc8;
-    font-family: 'JetBrains Mono', monospace;
+    font-size: 0.78rem;
+    color: #475569;
+    font-family: 'JetBrains Mono', monospace !important;
     padding: 0.35rem 0;
-    border-bottom: 1px solid #313244;
+    border-bottom: 1px solid #f1f5f9;
 }
 .dot {
     width: 8px; height: 8px;
     border-radius: 50%;
     flex-shrink: 0;
 }
-.dot-ok    { background: #a6e3a1; }
-.dot-fail  { background: #f38ba8; box-shadow: 0 0 6px #f38ba8; }
-.dot-final { background: #cba6f7; }
-
-.trace-summary { margin-top: 1rem; }
-.trace-label { font-size: 0.7rem; color: #a6adc8; text-transform: uppercase; letter-spacing: 0.08em; margin-bottom: 0.4rem; }
-.trace-bar-bg { background: #313244; border-radius: 4px; height: 6px; overflow: hidden; }
-.trace-bar { height: 100%; border-radius: 4px; transition: width 0.5s ease; }
-.trace-pct { font-size: 0.75rem; color: #a6adc8; margin-top: 0.3rem; font-family: 'JetBrains Mono', monospace; }
+.dot-ok    { background: #22c55e; }
+.dot-fail  { background: #ef4444; box-shadow: 0 0 5px rgba(239,68,68,0.4); }
+.dot-final { background: #7c3aed; }
 
-/* Panel labels */
-.panel-label {
-    font-size: 0.7rem;
-    font-weight: 600;
-    color: #a6adc8;
+.trace-summary { margin-top: 1rem; padding-top: 0.75rem; border-top: 1px solid #e2e8f0; }
+.trace-label {
+    font-size: 0.65rem;
+    font-weight: 700;
+    color: #64748b;
     text-transform: uppercase;
     letter-spacing: 0.1em;
-    margin-bottom: 0.75rem;
-    padding-bottom: 0.5rem;
-    border-bottom: 1px solid #313244;
+    margin-bottom: 0.5rem;
+}
+.trace-bar-bg {
+    background: #e2e8f0;
+    border-radius: 999px;
+    height: 6px;
+    overflow: hidden;
+}
+.trace-bar {
+    height: 100%;
+    border-radius: 999px;
+    transition: width 0.6s ease;
+}
+.trace-pct {
+    font-size: 0.72rem;
+    color: #475569;
+    margin-top: 0.3rem;
+    font-family: 'JetBrains Mono', monospace !important;
+    font-weight: 500;
 }
 
-/* Messages */
+/* ── Messages ── */
 .msg { padding: 0.75rem 1rem; border-radius: 8px; font-size: 0.85rem; }
-.msg.warn  { background: #2a2010; color: #f9e2af; border: 1px solid #f9e2af; }
-.msg.error { background: #2a1520; color: #f38ba8; border: 1px solid #f38ba8; }
-.msg.empty { color: #6c7086; font-size: 0.8rem; font-family: 'JetBrains Mono', monospace; }
-
-/* Output panels */
-.output-panel {
-    background: #181825;
-    border: 1px solid #313244;
-    border-radius: 10px;
+.msg.warn  { background: #fffbeb; color: #92400e; border: 1px solid #fcd34d; }
+.msg.error { background: #fff5f5; color: #b91c1c; border: 1px solid #fca5a5; }
+.msg.empty {
+    color: #94a3b8;
+    font-size: 0.8rem;
+    font-family: 'JetBrains Mono', monospace !important;
     padding: 1rem;
-    min-height: 200px;
 }
+
+/* ── Gradio overrides ── */
+.gradio-container .prose { color: #1e293b !important; }
+.gr-box, .gr-form { background: #ffffff !important; border-color: #e2e8f0 !important; }
+label span { color: #374151 !important; font-size: 0.85rem !important; }
 """
 
 # ── Layout ────────────────────────────────────────────────────────────────────
@@ -334,7 +358,7 @@ HEADER = """
 </div>
 """
 
-EXAMPLE_PROBLEMS = [
+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.",
@@ -349,20 +373,16 @@ with gr.Blocks(css=CSS, title="ThinkWhileThinking") as demo:
             gr.HTML("<div class='input-label'>Problem</div>")
             problem_input = gr.Textbox(
                 placeholder="Enter a math, logic, or reasoning problem...",
-                lines=4,
+                lines=5,
                 show_label=False,
             )
-            run_btn = gr.Button("▶ Analyze Reasoning", variant="primary")
-            gr.Examples(
-                examples=EXAMPLE_PROBLEMS,
-                inputs=problem_input,
-                label="Try an example",
-            )
+            run_btn = gr.Button("▶  Analyze Reasoning", variant="primary")
+            gr.Examples(examples=EXAMPLES, inputs=problem_input, label="Try an example")
 
         with gr.Column(scale=3):
             with gr.Row():
                 with gr.Column(scale=3):
-                    gr.HTML("<div class='panel-label'>Chain of Thought · Sentinel Scores</div>")
+                    gr.HTML("<div class='panel-label'>Chain of Thought · Step Scores</div>")
                     cot_output = gr.HTML(
                         value="<div class='msg empty'>// awaiting problem input</div>"
                     )
@@ -374,15 +394,7 @@ with gr.Blocks(css=CSS, title="ThinkWhileThinking") as demo:
 
     raw_output = gr.Textbox(visible=False)
 
-    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],
-    )
+    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()