Update app.py

app.py

@@ -102,6 +102,29 @@ GEN_TOP_K        = int(_gen_cfg.get('top_k',                1))
 GEN_REP_PENALTY  = float(_gen_cfg.get('repetition_penalty', 1.0))
 STEER_DISPLAY_K  = 10   # top-k candidates shown in the per-token probability panel
 
+# ─── Default chat templates (thinking / no-thinking) ─────────────────────────
+
+_THINK_TEMPLATE = (
+    "<|im_start|>user\n"
+    "{content}"
+    "<|im_end|>\n"
+    "<|im_start|>assistant\n"
+    "<think>\n"
+)
+
+_NOTHINK_TEMPLATE = (
+    "<|im_start|>user\n"
+    "{content}"
+    "<|im_end|>\n"
+    "<|im_start|>assistant\n"
+    "<think>\n\n</think>\n\n"
+)
+
+def apply_default_template(prompt: str, think: bool) -> str:
+    """Wrap *prompt* in the ChatML template for thinking or no-thinking mode."""
+    tpl = _THINK_TEMPLATE if think else _NOTHINK_TEMPLATE
+    return tpl.format(content=prompt.strip())
+
 # ─── Device resolution ───────────────────────────────────────────────────────
 
 def _resolve_sae_device() -> torch.device:
@@ -502,18 +525,19 @@ def _steering_strength_from_mode(mode: str, diff_lookup, layer: int, feat_idx: i
 
 def cb_generate(prompt, layer, feat_idx, pos_str, steer_mode, compare_diff,
                 steer_output_only, max_tok, greedy, top_k_tok, top_p, rep_penalty, temp,
-                custom_strength=5.0):
+                custom_strength=5.0, apply_think=False, apply_nothink=False):
     try:
         return _cb_generate_inner(prompt, layer, feat_idx, pos_str, steer_mode, compare_diff,
                                   steer_output_only, max_tok, greedy, top_k_tok, top_p, rep_penalty, temp,
-                                  custom_strength)
+                                  custom_strength, apply_think, apply_nothink)
     except gr.Error:
         raise
     except Exception as e:
         raise gr.Error(f"Generation failed: {e}")
 
 
-def cb_update_steer_preview(prompt: str, pos_str: str):
+def cb_update_steer_preview(prompt: str, pos_str: str,
+                             apply_think: bool = False, apply_nothink: bool = False):
     """Tokenise the prompt and return an HTML token-position preview."""
     if not prompt.strip():
         return (
@@ -522,7 +546,13 @@ def cb_update_steer_preview(prompt: str, pos_str: str):
         )
     try:
         _, tokenizer = get_model()
-        input_ids = tokenizer.encode(prompt)
+        if apply_think:
+            effective = apply_default_template(prompt, think=True)
+        elif apply_nothink:
+            effective = apply_default_template(prompt, think=False)
+        else:
+            effective = prompt
+        input_ids = tokenizer.encode(effective)
         tokens    = [tokenizer.decode([t]) for t in input_ids]
         positions = parse_positions(pos_str)
         return tokens_with_positions_html(tokens, positions)
@@ -706,7 +736,7 @@ def probs_to_html(tokens: list, chosen_probs: list, topk_data: list,
 
 def _cb_generate_inner(prompt, layer, feat_idx, pos_str, steer_mode, compare_diff,
                        steer_output_only, max_tok, greedy, top_k_tok, top_p, rep_penalty, temp,
-                       custom_strength=5.0):
+                       custom_strength=5.0, apply_think=False, apply_nothink=False):
     global _orig_cache
     model, tokenizer = get_model()
     layer    = int(layer)
@@ -716,7 +746,13 @@ def _cb_generate_inner(prompt, layer, feat_idx, pos_str, steer_mode, compare_dif
     strength  = _steering_strength_from_mode(steer_mode, compare_diff, layer, feat_idx, float(custom_strength))
     positions = parse_positions(pos_str)
 
-    input_ids = tokenizer.encode(prompt, return_tensors='pt').to(
+    if apply_think:
+        effective_prompt = apply_default_template(prompt, think=True)
+    elif apply_nothink:
+        effective_prompt = apply_default_template(prompt, think=False)
+    else:
+        effective_prompt = prompt
+    input_ids = tokenizer.encode(effective_prompt, return_tensors='pt').to(
         next(model.parameters()).device
     )
 
@@ -738,9 +774,9 @@ def _cb_generate_inner(prompt, layer, feat_idx, pos_str, steer_mode, compare_dif
     # The unsteered output depends only on the prompt and decoding parameters,
     # not on any steering inputs.  Reuse the last result when those are unchanged.
     if greedy:
-        orig_key = (prompt, int(max_tok), True)
+        orig_key = (effective_prompt, int(max_tok), True)
     else:
-        orig_key = (prompt, int(max_tok), False,
+        orig_key = (effective_prompt, int(max_tok), False,
                     int(top_k_tok), float(top_p), float(rep_penalty), float(temp))
 
     if _orig_cache is not None and _orig_cache['key'] == orig_key:
@@ -1495,6 +1531,24 @@ with gr.Blocks(title="Qwen-Scope Feature Explorer") as demo:
                             show_label=False,
                         )
 
+                        t2_apply_think = gr.Checkbox(
+                            label="Apply default thinking template",
+                            value=False,
+                            info=(
+                                "Wrap the prompt in the ChatML format with thinking enabled "
+                                "(assistant prefill starts with <think>)."
+                            ),
+                        )
+                        t2_apply_nothink = gr.Checkbox(
+                            label="Apply default no-thinking template",
+                            value=False,
+                            info=(
+                                "Wrap the prompt in the ChatML format with thinking disabled "
+                                "(assistant prefill starts with <think>\\n\\n</think>)."
+                            ),
+                        )
+                        t2_template_info = gr.HTML(visible=False, value="")
+
                         gr.HTML('<span class="section-chip">Token Position Preview</span>'
                                 '<span style="font-size:12px;color:#888;margin-left:8px;">'
                                 'amber = steered &nbsp;·&nbsp; updates as you type'
@@ -1676,7 +1730,7 @@ with gr.Blocks(title="Qwen-Scope Feature Explorer") as demo:
                 inputs=[t2_prompt, t2_layer, t2_feat, t2_pos, t2_steer_mode, compare_diff_state,
                         t2_steer_output_only, t2_maxtok,
                         t2_greedy, t2_top_k_tok, t2_top_p, t2_rep_penalty,
-                        t2_temperature, t2_custom_strength],
+                        t2_temperature, t2_custom_strength, t2_apply_think, t2_apply_nothink],
                 outputs=[t2_orig, t2_steered, t2_orig_probs, t2_steer_probs],
             )
             t3_run.click(
@@ -1699,12 +1753,12 @@ with gr.Blocks(title="Qwen-Scope Feature Explorer") as demo:
             )
             t2_prompt.change(
                 cb_update_steer_preview,
-                inputs=[t2_prompt, t2_pos],
+                inputs=[t2_prompt, t2_pos, t2_apply_think, t2_apply_nothink],
                 outputs=[t2_pos_preview],
             )
             t2_pos.change(
                 cb_update_steer_preview,
-                inputs=[t2_prompt, t2_pos],
+                inputs=[t2_prompt, t2_pos, t2_apply_think, t2_apply_nothink],
                 outputs=[t2_pos_preview],
             )
 
@@ -1763,6 +1817,69 @@ with gr.Blocks(title="Qwen-Scope Feature Explorer") as demo:
                 outputs=[t2_custom_strength],
             )
 
+            # ── Template toggle: mutual exclusion + info panel + preview refresh ─
+            _THINK_INFO_HTML = (
+                '<div style="font-size:11px;color:#555;padding:6px 10px;'
+                'background:#eff6ff;border:1px solid #bfdbfe;border-radius:6px;'
+                'font-family:ui-monospace,monospace;white-space:pre-wrap;line-height:1.7;">'
+                '&lt;|im_start|&gt;user\n'
+                '&#123;your prompt&#125;&lt;|im_end|&gt;\n'
+                '&lt;|im_start|&gt;assistant\n'
+                '&lt;think&gt;\n'
+                '</div>'
+            )
+            _NOTHINK_INFO_HTML = (
+                '<div style="font-size:11px;color:#555;padding:6px 10px;'
+                'background:#f0fdf4;border:1px solid #bbf7d0;border-radius:6px;'
+                'font-family:ui-monospace,monospace;white-space:pre-wrap;line-height:1.7;">'
+                '&lt;|im_start|&gt;user\n'
+                '&#123;your prompt&#125;&lt;|im_end|&gt;\n'
+                '&lt;|im_start|&gt;assistant\n'
+                '&lt;think&gt;\n\n&lt;/think&gt;\n\n'
+                '</div>'
+            )
+
+            def _on_think_change(think_val, nothink_val, prompt, pos_str):
+                if think_val:
+                    # Just checked: uncheck nothink, show think format, refresh preview
+                    return (gr.update(value=False),
+                            gr.update(visible=True, value=_THINK_INFO_HTML),
+                            cb_update_steer_preview(prompt, pos_str, True, False))
+                elif nothink_val:
+                    # Unchecked by mutual exclusion — nothink is active; leave info+preview alone
+                    return gr.update(), gr.update(), gr.update()
+                else:
+                    # Manually unchecked with nothing active — reset to raw
+                    return (gr.update(),
+                            gr.update(visible=False),
+                            cb_update_steer_preview(prompt, pos_str, False, False))
+
+            def _on_nothink_change(nothink_val, think_val, prompt, pos_str):
+                if nothink_val:
+                    # Just checked: uncheck think, show nothink format, refresh preview
+                    return (gr.update(value=False),
+                            gr.update(visible=True, value=_NOTHINK_INFO_HTML),
+                            cb_update_steer_preview(prompt, pos_str, False, True))
+                elif think_val:
+                    # Unchecked by mutual exclusion — think is active; leave info+preview alone
+                    return gr.update(), gr.update(), gr.update()
+                else:
+                    # Manually unchecked with nothing active — reset to raw
+                    return (gr.update(),
+                            gr.update(visible=False),
+                            cb_update_steer_preview(prompt, pos_str, False, False))
+
+            t2_apply_think.change(
+                fn=_on_think_change,
+                inputs=[t2_apply_think, t2_apply_nothink, t2_prompt, t2_pos],
+                outputs=[t2_apply_nothink, t2_template_info, t2_pos_preview],
+            )
+            t2_apply_nothink.change(
+                fn=_on_nothink_change,
+                inputs=[t2_apply_nothink, t2_apply_think, t2_prompt, t2_pos],
+                outputs=[t2_apply_think, t2_template_info, t2_pos_preview],
+            )
+
 
 if __name__ == '__main__':
     # Pre-load model onto GPU before accepting requests, so the first