o1 impl

app.py

@@ -109,7 +109,7 @@ def generate(
         "unit_locations": None,
         "max_new_tokens": max_new_tokens,
         "intervene_on_prompt": True,
-        "subspaces": [{"idx": subspaces_list["idx"], "mag": subspaces_list["internal_mag"]}],
+        "subspaces": [{"idx": int(subspaces_list[0]["idx"]), "mag": int(subspaces_list[0]["internal_mag"])}],
         "streamer": streamer,
         "do_sample": True
     }
@@ -185,13 +185,13 @@ with gr.Blocks(css="style.css") as demo:
         with gr.Column(scale=7):
             chat_interface = gr.ChatInterface(
                 fn=generate,
-                title="## LM Steering with ReFT-r1 (16K concepts)",
+                title="LM Steering with ReFT-r1 (16K concepts)",
                 type="messages",
                 additional_inputs=[selected_subspaces],
             )
         # Right side: concept management
         with gr.Column(scale=3):
-            gr.Markdown("### Steering Concepts")
+            gr.Markdown("## Steering Concepts")
             search_box = gr.Textbox(
                 label="Search concepts",
                 placeholder="e.g. 'time travel'"