update

.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+dummy/

app.py

@@ -98,10 +98,6 @@ div#banner {
   
 }
 
-div#main-components {
-    align-items: flex-end;
-}
-
 div#steering-toggle {
     padding-top: 8px;
     padding-bottom: 8px;
@@ -233,13 +229,31 @@ async def get_endpoint_state():
         yield "Server Error"
 
 
+async def post_process(session_id, output):
+    req = instances[session_id].pop()
+
+    if "</think>" in output:
+        p = [p for p in output.partition("</think>") if p != ""]
+        reasoning = "".join(p[:-1])
+        if len(p) == 1:
+            answer = None
+        else:
+            answer = p[-1]
+    else:
+        answer = None
+        reasoning = output
+
+    steering_output = SteeringOutput(**req.model_dump(), reasoning=reasoning, answer=answer)
+    instances[session_id].append(steering_output)
+
+
 async def generate(
     session_id: str, prompt: str, steering: bool, coeff: float, 
-    max_new_tokens: int, top_p: float, temperature: float
+    max_new_tokens: int, top_p: float, temperature: float, 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
+        max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature, k=vec_scaling
     )
 
     instances[session_id].append(req)
@@ -260,28 +274,11 @@ async def generate(
                 else:
                     logger.error("API Error Ccode: %d, Error Message: %s", resp.status, resp.text())
                     raise gr.Error("API Server Error")
-                
-    except:
-        logger.info("Client session Error")
-
-
-async def post_process(session_id, output):
-    req = instances[session_id].pop()
-
-    if "</think>" in output:
-        p = [p for p in output.partition("</think>") if p != ""]
-        reasoning = "".join(p[:-1])
-        if len(p) == 1:
-            answer = None
-        else:
-            answer = p[-1]
-    else:
-        answer = None
-        reasoning = output
+        
+        await post_process(session_id, generated_text)
 
-    steering_output = SteeringOutput(**req.model_dump(), reasoning=reasoning, answer=answer)
-    instances[session_id].append(steering_output)
-    return gr.update(interactive=True), gr.update(interactive=True)
+    except:
+        logger.info("Client session error")
 
 
 async def output_feedback(session_id, feedback):
@@ -299,6 +296,10 @@ async def output_feedback(session_id, feedback):
         logger.debug("Feedback submission error")
 
 
+async def show_feedback_buttons(upvote_btn, downvote_btn):
+    return gr.update(interactive=True), gr.update(interactive=True)
+
+
 gr.set_static_paths(paths=[Path.cwd().absolute() / "assets"])
 theme = gr.themes.Base(primary_hue="emerald", text_size=gr.themes.sizes.text_lg).set()
 
@@ -308,20 +309,20 @@ with gr.Blocks(title="LLM Censorship Steering", theme=theme, head=HEAD, css=CSS,
 
     gr.HTML(HTML)
 
+    @gr.render(inputs=endpoint_state, triggers=[endpoint_state.change])
+    def render_state(endpoint_state):
+        if endpoint_state == "Ready":
+            color = "green"
+        elif endpoint_state == "Server Error":
+            color = "red"
+        else:
+            color = "orange"
+
+        if endpoint_state != None:
+            gr.Markdown(f'🤖 {model_name} | Inference Endpoint State: <span style="color:{color}; font-weight: bold;">{endpoint_state}</span>', elem_id="model-state")
+
     with gr.Row(elem_id="main-components"):
         with gr.Column(scale=1):
-            @gr.render(inputs=endpoint_state, triggers=[endpoint_state.change])
-            def render_state(endpoint_state):
-                if endpoint_state == "Ready":
-                    color = "green"
-                elif endpoint_state == "Server Error":
-                    color = "red"
-                else:
-                    color = "orange"
-
-                if endpoint_state != None:
-                    gr.Markdown(f'🤖 {model_name} | Inference Endpoint State: <span style="color:{color}; font-weight: bold;">{endpoint_state}</span>', elem_id="model-state")
-                    
             with gr.Row():
                 steer_toggle = Toggle(label="Steering", info="Turn off to generate original outputs", value=True, interactive=True, scale=2, elem_id="steering-toggle")
                 coeff = gr.Slider(label="Coefficient:", value=-1.0, minimum=-2, maximum=2, step=0.1, scale=8, show_reset_button=False, elem_id="coeff-slider")
@@ -332,23 +333,28 @@ with gr.Blocks(title="LLM Censorship Steering", theme=theme, head=HEAD, css=CSS,
                     return gr.update(label="Steering", info="Turn off to generate original outputs"), gr.update(interactive=True)
                 else:
                     return gr.update(label="No Steering", info="Turn on to steer model outputs"), gr.update(interactive=False)
-
-            with gr.Accordion("⚙️ Advanced Settings", open=False):
-                with gr.Row():
-                    temperature = gr.Slider(0, 1, step=0.1, value=CONFIG["temperature"], interactive=True, label="Temperature", scale=2)
-                    top_p = gr.Slider(0, 1, step=0.1, value=CONFIG["top_p"], interactive=True, label="Top p", scale=2)
-                    max_new_tokens = gr.Number(CONFIG["max_new_tokens"], minimum=10, maximum=CONFIG["max_new_tokens"], interactive=True, label="Max new tokens", scale=1)
             
             input_text = gr.Textbox(label="Input", placeholder="Enter your prompt here...", lines=6, interactive=True)
 
             with gr.Row():
                 clear_btn = gr.ClearButton()
+                stop_btn = gr.Button("Stop")
                 generate_btn = gr.Button("Generate", variant="primary")
 
+            with gr.Accordion("⚙️ Advanced Settings", open=False):
+                with gr.Row():
+                    temperature = gr.Slider(0, 1, step=0.1, value=CONFIG["temperature"], interactive=True, label="Temperature", scale=1)
+                    top_p = gr.Slider(0, 1, step=0.1, value=CONFIG["top_p"], interactive=True, label="Top p", scale=1)
+
+                with gr.Row():
+                    layer = gr.Slider(0, 27, step=1, value=CONFIG["layer"], interactive=True, label="Steering layer", scale=2)
+                    max_new_tokens = gr.Number(CONFIG["max_new_tokens"], minimum=10, maximum=CONFIG["max_new_tokens"], interactive=True, label="Max new tokens", scale=1)
+                    vec_scaling = gr.Number(CONFIG["k"], interactive=True, label="Vector scaling", scale=1)
+
         with gr.Column(scale=1):
             output = gr.Textbox(label="Output", lines=15, max_lines=15, interactive=False)
 
-            with gr.Row():
+            with gr.Row():                
                 upvote_btn = gr.Button("👍 Upvote", interactive=False)
                 downvote_btn = gr.Button("👎 Downvote", interactive=False)
 
@@ -357,17 +363,23 @@ with gr.Blocks(title="LLM Censorship Steering", theme=theme, head=HEAD, css=CSS,
     gr.Examples(examples=examples[examples["type"] == "sensitive"].prompt.tolist(), inputs=input_text, label="Sensitive")
     gr.Examples(examples=examples[examples["type"] == "harmful"].prompt.tolist(), inputs=input_text, label="Harmful")
 
-    @gr.on(triggers=[clear_btn.click], outputs=[upvote_btn, downvote_btn])
-    def clear():
+
+    @gr.on(triggers=[clear_btn.click, stop_btn.click], outputs=[upvote_btn, downvote_btn])
+    def clear_feedback_buttons():
         return gr.update(interactive=False), gr.update(interactive=False)
+    
+    @gr.on(triggers=[generate_btn.click], outputs=[upvote_btn, downvote_btn])
+    def show_feedback_buttons():
+        return gr.update(interactive=True), gr.update(interactive=True)
 
-    clear_btn.add([input_text, output])
-    generate_btn.click(
-        generate, inputs=[session_id, input_text, steer_toggle, coeff, max_new_tokens, top_p, temperature], outputs=output
-    ).success(
-        post_process, inputs=[session_id, output], outputs=[upvote_btn, downvote_btn]
+
+    submission = generate_btn.click(
+        generate, inputs=[session_id, input_text, steer_toggle, coeff, max_new_tokens, top_p, temperature, vec_scaling], outputs=output
     )
 
+    clear_btn.add([input_text, output])
+    stop_btn.click(None, None, None, cancels=[submission], queue=False)
+
     upvote_btn.click(output_feedback, inputs=[session_id, upvote_btn])
     downvote_btn.click(output_feedback, inputs=[session_id, downvote_btn])
 

scheduler.py

@@ -28,6 +28,8 @@ def load_scheduler():
             "answer": {"_type": "Value", "dtype": "string"},
             "upvote": {"_type": "Value", "dtype": "bool"},
             "timestamp": {"_type": "Value", "dtype": "string"},
+            "layer": {"_type": "Value", "dtype": "int64"}, 
+            "k": {"_type": "Value", "dtype": "float64"},
         }
     )
 

schemas.py

@@ -7,7 +7,8 @@ CONFIG = {
     "max_new_tokens": 3048,
     "top_p": 0.95,
     "temperature": 0.6,
-    "k": 200
+    "k": 200,
+    "layer": 25
 }
 
 class UserRequest(BaseModel):
@@ -18,13 +19,16 @@ class UserRequest(BaseModel):
     max_new_tokens: int = Field(CONFIG["max_new_tokens"], le=3048)
     top_p: float = Field(CONFIG["top_p"], ge=0.0, le=1.0)
     temperature: float = Field(CONFIG["temperature"], ge=0.0, le=1.0)
+    k: float = Field(CONFIG["k"])
+    layer: int = Field(CONFIG["layer"])
 
     def get_api_format(self):
         return {
             "prompt": self.prompt,
             "steering": self.steering,
             "coeff": self.coeff,
-            "k": CONFIG["k"],
+            "k": self.k,
+            "layer": self.layer,
             "generation_config": {
                 "max_new_tokens": self.max_new_tokens,
                 "top_p": self.top_p,
@@ -36,6 +40,6 @@ class UserRequest(BaseModel):
 class SteeringOutput(UserRequest):
     max_new_tokens: SkipJsonSchema[int] = Field(exclude=True)
     reasoning: str = None
-    answer: str = None
+    answer: Optional[str] = None
     upvote: Optional[bool] = None
     timestamp: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat())