| import gradio as gr |
| import time |
| import spaces |
| import torch |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForCausalLM, pipeline |
|
|
| |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" |
| print(f"🖥️ Device: {DEVICE}") |
|
|
| |
| REASONER_MODEL = "Qwen/Qwen2.5-1.5B-Instruct" |
| print(f"🔄 Loading reasoner: {REASONER_MODEL}") |
| reasoner_tokenizer = AutoTokenizer.from_pretrained(REASONER_MODEL, trust_remote_code=True) |
| reasoner_model = AutoModelForCausalLM.from_pretrained( |
| REASONER_MODEL, |
| trust_remote_code=True, |
| torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32, |
| device_map="auto", |
| ) |
| reasoner_pipe = pipeline( |
| "text-generation", |
| model=reasoner_model, |
| tokenizer=reasoner_tokenizer, |
| max_new_tokens=600, |
| temperature=0.3, |
| do_sample=True, |
| ) |
| print("✅ Reasoner loaded!") |
|
|
| |
| PROBE_MODEL = "realArceus/twt-probe" |
| print(f"🔄 Loading probe: {PROBE_MODEL}") |
| probe_tokenizer = AutoTokenizer.from_pretrained(PROBE_MODEL, trust_remote_code=True) |
| probe_model = AutoModelForSequenceClassification.from_pretrained( |
| PROBE_MODEL, |
| trust_remote_code=True, |
| torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32, |
| ).to(DEVICE) |
| probe_model.eval() |
| print("✅ Probe loaded!") |
|
|
| SYSTEM_PROMPT = """You are a careful step-by-step reasoner. When given a problem, solve it by thinking through exactly numbered steps. |
| Format EVERY step as: |
| Step 1: <your reasoning> |
| Step 2: <your reasoning> |
| ... |
| Final Answer: <answer> |
| Be deliberate. Show your full working. Each step should be one clear thought.""" |
|
|
| |
| def probe_step(problem: str, steps_so_far: list) -> tuple: |
| context = " ".join(steps_so_far[:-1]) if len(steps_so_far) > 1 else "" |
| current = steps_so_far[-1] if steps_so_far else "" |
| input_text = ( |
| f"[PROBLEM] {problem.strip()} " |
| f"[STEPS SO FAR] {context} " |
| f"[CURRENT STEP] {current}" |
| ) |
| inputs = probe_tokenizer( |
| input_text, return_tensors="pt", |
| truncation=True, max_length=512, |
| ).to(DEVICE) |
| with torch.no_grad(): |
| logits = probe_model(**inputs).logits |
| probs = torch.softmax(logits, dim=-1)[0] |
| ok_conf = probs[1].item() |
| fail_conf = probs[0].item() |
| label = "OK" if ok_conf >= 0.5 else "FAIL" |
| conf = round(ok_conf if label == "OK" else fail_conf, 2) |
| return label, conf |
|
|
| |
| @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>") |
| return |
|
|
| cot_html = "" |
|
|
| messages = [ |
| {"role": "system", "content": SYSTEM_PROMPT}, |
| {"role": "user", "content": problem.strip()} |
| ] |
| prompt = reasoner_tokenizer.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True |
| ) |
|
|
| try: |
| result = reasoner_pipe(prompt) |
| full_response = result[0]["generated_text"][len(prompt):] |
| except Exception as e: |
| yield ("", f"<div class='msg error'>❌ Reasoner error: {str(e)}</div>", "") |
| return |
|
|
| lines = full_response.strip().split("\n") |
| parsed_steps = [l.strip() for l in lines if l.strip()] |
| displayed_steps = [] |
| first_fail_seen = False |
|
|
| 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) |
| else: |
| label, conf = "FINAL", 1.0 |
|
|
| if is_final: |
| icon, badge_class, badge_text = "🏁", "badge-final", "ANSWER" |
| elif label == "OK": |
| icon, badge_class, badge_text = "🟢", "badge-ok", f"OK · {int(conf*100)}%" |
| else: |
| icon, badge_class, badge_text = "🔴", "badge-fail", f"FAULT · {int(conf*100)}%" |
| if not first_fail_seen: |
| first_fail_seen = True |
|
|
| step_card = f""" |
| <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 |
|
|
| yield (full_response, cot_html, build_trace(problem, displayed_steps, first_fail_seen)) |
| time.sleep(0.1) |
|
|
| yield (full_response, cot_html, build_trace(problem, displayed_steps, first_fail_seen, done=True)) |
|
|
|
|
| def build_trace(problem, displayed, first_fail, done=False): |
| faults, rows = 0, "" |
| 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]) |
| 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'>— {int(conf*100)}%</span></div>" |
| else: |
| rows += "<div class='trace-row'><span class='dot dot-final'></span> Final Answer</div>" |
|
|
| health = max(0, 100 - faults * 20) |
| bar_color = "#10b981" if health > 60 else "#e11d48" |
| summary = f""" |
| <div class='trace-summary'> |
| <div class='trace-label'>Reasoning Health</div> |
| <div class='trace-bar-bg'><div class='trace-bar' style='width:{health}%;background:{bar_color}'></div></div> |
| <div class='trace-pct'>{health}%</div> |
| </div>""" if done else "" |
| return f"{rows}{summary}" |
|
|
|
|
| |
| CSS = """ |
| @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=JetBrains+Mono:wght@400;500;600&display=swap'); |
| |
| *, *::before, *::after { box-sizing: border-box; } |
| body, .gradio-container, .gradio-container * { |
| font-family: 'Inter', -apple-system, sans-serif !important; |
| color: #1e293b !important; |
| } |
| .gradio-container { |
| background-color: #f4f6f8 !important; |
| max-width: 1280px !important; |
| padding-top: 1rem !important; |
| } |
| .twt-header { |
| text-align: center; |
| padding: 0 1rem 1rem; |
| margin-bottom: 1.5rem; |
| border-bottom: 1px solid #cbd5e1; |
| } |
| .twt-title { |
| 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; } |
| .twt-sub { |
| font-size: 0.85rem; |
| color: #475569 !important; |
| font-weight: 500; |
| letter-spacing: 0.02em; |
| } |
| .input-label, .panel-label { |
| font-size: 0.7rem; |
| font-weight: 700; |
| color: #334155 !important; |
| text-transform: uppercase; |
| letter-spacing: 0.08em; |
| margin-bottom: 0.5rem; |
| display: block; |
| } |
| .panel-label { |
| padding-bottom: 0.4rem; |
| border-bottom: 2px solid #cbd5e1; |
| } |
| textarea { |
| background: #ffffff !important; |
| border: 1px solid #94a3b8 !important; |
| border-radius: 6px !important; |
| color: #0f172a !important; |
| font-size: 0.85rem !important; |
| line-height: 1.5 !important; |
| padding: 0.75rem !important; |
| box-shadow: inset 0 1px 2px rgba(15, 23, 42, 0.05) !important; |
| transition: all 0.2s ease !important; |
| } |
| textarea:focus { |
| border-color: #0033a0 !important; |
| box-shadow: 0 0 0 3px rgba(0, 51, 160, 0.15) !important; |
| outline: none !important; |
| } |
| textarea::placeholder { color: #94a3b8 !important; } |
| button.primary { |
| background: #0033a0 !important; |
| color: #ffffff !important; |
| font-weight: 600 !important; |
| border: none !important; |
| 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; |
| width: 100% !important; |
| } |
| button.primary:hover { |
| background: #002266 !important; |
| transform: translateY(-1px) !important; |
| box-shadow: 0 4px 6px -2px rgba(0, 51, 160, 0.3) !important; |
| } |
| button.primary:active { |
| transform: translateY(0) !important; |
| box-shadow: 0 1px 2px rgba(0, 51, 160, 0.2) !important; |
| } |
| .step-card { |
| display: flex; |
| align-items: flex-start; |
| gap: 0.75rem; |
| padding: 0.75rem 1rem; |
| margin-bottom: 0.5rem; |
| background: #ffffff; |
| border: 1px solid #cbd5e1; |
| border-radius: 6px; |
| line-height: 1.5; |
| animation: fadeUp 0.2s ease-out both; |
| box-shadow: 0 1px 2px rgba(15, 23, 42, 0.03); |
| transition: all 0.15s ease; |
| } |
| .step-card:hover { |
| box-shadow: 0 2px 6px rgba(15, 23, 42, 0.06); |
| border-color: #94a3b8; |
| } |
| .step-card.step-fault { |
| border-color: #fca5a5; |
| background: #fff1f2; |
| } |
| .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; } |
| .step-text { flex: 1; color: #1e293b !important; font-size: 0.85rem; } |
| .badge { |
| flex-shrink: 0; |
| font-size: 0.6rem; |
| font-weight: 700; |
| font-family: 'JetBrains Mono', monospace !important; |
| letter-spacing: 0.02em; |
| padding: 0.15rem 0.4rem; |
| border-radius: 4px; |
| margin-top: 2px; |
| text-transform: uppercase; |
| } |
| .badge-ok { background: #f0fdf4; color: #15803d !important; border: 1px solid #86efac; } |
| .badge-fail { background: #fff1f2; color: #be123c !important; border: 1px solid #fda4af; } |
| .badge-final { background: #f0f4ff; color: #0033a0 !important; border: 1px solid #bfdbfe; } |
| .trace-row { |
| display: flex; |
| align-items: center; |
| gap: 0.5rem; |
| font-size: 0.75rem; |
| color: #334155 !important; |
| font-family: 'JetBrains Mono', monospace !important; |
| 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; } |
| .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; } |
| .trace-summary { |
| margin-top: 1rem; |
| padding-top: 0.75rem; |
| border-top: 1px solid #cbd5e1; |
| } |
| .trace-label { |
| font-size: 0.65rem; |
| font-weight: 700; |
| color: #475569 !important; |
| text-transform: uppercase; |
| letter-spacing: 0.05em; |
| margin-bottom: 0.4rem; |
| } |
| .trace-bar-bg { |
| background: #cbd5e1; |
| border-radius: 999px; |
| height: 4px; |
| overflow: hidden; |
| } |
| .trace-bar { |
| height: 100%; |
| border-radius: 999px; |
| transition: width 0.8s cubic-bezier(0.16, 1, 0.3, 1); |
| } |
| .trace-pct { |
| font-size: 0.75rem; |
| color: #0f172a !important; |
| margin-top: 0.4rem; |
| font-family: 'JetBrains Mono', monospace !important; |
| font-weight: 600; |
| } |
| .msg { padding: 0.75rem; border-radius: 6px; font-size: 0.8rem; font-weight: 500; } |
| .msg.warn { background: #fefce8; color: #a16207 !important; border: 1px solid #fde047; } |
| .msg.error { background: #fff1f2; color: #be123c !important; border: 1px solid #fca5a5; } |
| .msg.empty { |
| color: #64748b !important; |
| font-size: 0.8rem; |
| font-family: 'JetBrains Mono', monospace !important; |
| padding: 1.5rem 1rem; |
| text-align: center; |
| background: #ffffff; |
| border: 1px dashed #94a3b8; |
| border-radius: 6px; |
| } |
| .gr-box, .gr-form { background: transparent !important; border: none !important; } |
| """ |
|
|
| HEADER = """ |
| <div class='twt-header'> |
| <div class='twt-title'>Think<span>While</span>Thinking</div> |
| <div class='twt-sub'>Real-time reasoning failure detection · Step-level process supervision</div> |
| </div> |
| """ |
|
|
| 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?", |
| ] |
|
|
| with gr.Blocks(css=CSS, title="ThinkWhileThinking") as demo: |
| gr.HTML(HEADER) |
| with gr.Row(): |
| with gr.Column(scale=2): |
| gr.HTML("<div class='input-label'>Problem</div>") |
| problem_input = gr.Textbox(placeholder="Enter a math, logic, or reasoning problem...", lines=5, show_label=False) |
| 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 · Step Scores</div>") |
| cot_output = gr.HTML(value="<div class='msg empty'>// awaiting problem input</div>") |
| with gr.Column(scale=1): |
| gr.HTML("<div class='panel-label'>Step Trace</div>") |
| trace_output = gr.HTML(value="<div class='msg empty'>// trace</div>") |
| 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]) |
|
|
| demo.launch() |
|
|
