o1 impl

app.py

@@ -104,13 +104,12 @@ def generate(
         yield "[Truncated prior text]\n"
 
     streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-    print(subspaces_list)
     generate_kwargs = {
         "base": {"input_ids": input_ids},
         "unit_locations": None,
         "max_new_tokens": max_new_tokens,
         "intervene_on_prompt": True,
-        "subspaces": subspaces_list,
+        "subspaces": [{"idx": subspaces_list["idx"], "mag": subspaces_list["internal_mag"]}],
         "streamer": streamer,
         "do_sample": True
     }
@@ -168,8 +167,6 @@ def update_dropdown_choices(search_text):
     return gr.update(choices=filtered)
 
 with gr.Blocks(css="style.css") as demo:
-    # A short title only
-    gr.Markdown("## Model Steering with ReFT-r1 (16K concepts)")
 
     # Pre-populate with a random concept if available
     default_subspaces = []
@@ -188,7 +185,7 @@ with gr.Blocks(css="style.css") as demo:
         with gr.Column(scale=7):
             chat_interface = gr.ChatInterface(
                 fn=generate,
-                title="",
+                title="## LM Steering with ReFT-r1 (16K concepts)",
                 type="messages",
                 additional_inputs=[selected_subspaces],
             )