update

app.py

@@ -249,11 +249,11 @@ async def post_process(session_id, output):
 
 async def generate(
     session_id: str, prompt: str, steering: bool, coeff: float, 
-    max_new_tokens: int, top_p: float, temperature: float, vec_scaling: float
+    max_new_tokens: int, top_p: float, temperature: float, layer: int, vec_scaling: float
 ):
     req = UserRequest(
         session_id=session_id, prompt=prompt, steering=steering, coeff=coeff,
-        max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature, k=vec_scaling
+        max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature, layer=layer, k=vec_scaling
     )
 
     instances[session_id].append(req)
@@ -374,7 +374,7 @@ with gr.Blocks(title="LLM Censorship Steering", theme=theme, head=HEAD, css=CSS,
 
 
     submission = generate_btn.click(
-        generate, inputs=[session_id, input_text, steer_toggle, coeff, max_new_tokens, top_p, temperature, vec_scaling], outputs=output
+        generate, inputs=[session_id, input_text, steer_toggle, coeff, max_new_tokens, top_p, temperature, layer, vec_scaling], outputs=output
     )
 
     clear_btn.add([input_text, output])
@@ -383,6 +383,8 @@ with gr.Blocks(title="LLM Censorship Steering", theme=theme, head=HEAD, css=CSS,
     upvote_btn.click(output_feedback, inputs=[session_id, upvote_btn])
     downvote_btn.click(output_feedback, inputs=[session_id, downvote_btn])
 
+    layer.change(fn=lambda x: 1, inputs=vec_scaling, outputs=vec_scaling)
+
     demo.load(initialize_instance, outputs=session_id)
     demo.unload(cleanup_instance)