update

app.py

@@ -349,7 +349,7 @@ with gr.Blocks(title="LLM Censorship Steering", theme=theme, head=HEAD, css=CSS,
                 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)
+                    vec_scaling = gr.Number(1.0, minimum=0, interactive=True, label="Vector scaling", scale=1)
 
         with gr.Column(scale=1):
             output = gr.Textbox(label="Output", lines=15, max_lines=15, interactive=False)

assets/vector_scaling.json

@@ -0,0 +1,142 @@
+[
+    {
+        "layer": 0,
+        "k_pos": 4.0,
+        "k_neg": 4.2
+    },
+    {
+        "layer": 1,
+        "k_pos": 5.1,
+        "k_neg": 5.4
+    },
+    {
+        "layer": 2,
+        "k_pos": 8.2,
+        "k_neg": 7.9
+    },
+    {
+        "layer": 3,
+        "k_pos": 12.1,
+        "k_neg": 12.3
+    },
+    {
+        "layer": 4,
+        "k_pos": 13.4,
+        "k_neg": 13.9
+    },
+    {
+        "layer": 5,
+        "k_pos": 13.3,
+        "k_neg": 15.1
+    },
+    {
+        "layer": 6,
+        "k_pos": 16.3,
+        "k_neg": 17.2
+    },
+    {
+        "layer": 7,
+        "k_pos": 20.6,
+        "k_neg": 20.9
+    },
+    {
+        "layer": 8,
+        "k_pos": 28.9,
+        "k_neg": 28.3
+    },
+    {
+        "layer": 9,
+        "k_pos": 41.7,
+        "k_neg": 34.4
+    },
+    {
+        "layer": 10,
+        "k_pos": 43.7,
+        "k_neg": 35.1
+    },
+    {
+        "layer": 11,
+        "k_pos": 44.1,
+        "k_neg": 37.1
+    },
+    {
+        "layer": 12,
+        "k_pos": 48.9,
+        "k_neg": 41.1
+    },
+    {
+        "layer": 13,
+        "k_pos": 52.2,
+        "k_neg": 44.4
+    },
+    {
+        "layer": 14,
+        "k_pos": 56.0,
+        "k_neg": 48.3
+    },
+    {
+        "layer": 15,
+        "k_pos": 59.2,
+        "k_neg": 50.1
+    },
+    {
+        "layer": 16,
+        "k_pos": 63.0,
+        "k_neg": 52.6
+    },
+    {
+        "layer": 17,
+        "k_pos": 68.9,
+        "k_neg": 56.2
+    },
+    {
+        "layer": 18,
+        "k_pos": 76.2,
+        "k_neg": 63.4
+    },
+    {
+        "layer": 19,
+        "k_pos": 85.6,
+        "k_neg": 70.3
+    },
+    {
+        "layer": 20,
+        "k_pos": 96.7,
+        "k_neg": 78.9
+    },
+    {
+        "layer": 21,
+        "k_pos": 117.2,
+        "k_neg": 86.8
+    },
+    {
+        "layer": 22,
+        "k_pos": 135.0,
+        "k_neg": 99.8
+    },
+    {
+        "layer": 23,
+        "k_pos": 151.7,
+        "k_neg": 110.1
+    },
+    {
+        "layer": 24,
+        "k_pos": 172.6,
+        "k_neg": 125.6
+    },
+    {
+        "layer": 25,
+        "k_pos": 193.7,
+        "k_neg": 148.4
+    },
+    {
+        "layer": 26,
+        "k_pos": 217.1,
+        "k_neg": 164.0
+    },
+    {
+        "layer": 27,
+        "k_pos": 238.8,
+        "k_neg": 197.9
+    }
+]

schemas.py

@@ -1,13 +1,16 @@
+import json
 from typing import Optional
 from datetime import datetime, timezone
 from pydantic import BaseModel, Field
 from pydantic.json_schema import SkipJsonSchema
 
+vector_scaling = json.load(open("assets/vector_scaling.json", "r"))
+
 CONFIG = {
     "max_new_tokens": 3048,
     "top_p": 0.95,
     "temperature": 0.6,
-    "k": 200,
+    "k": vector_scaling[25]["k_pos"],
     "layer": 25
 }
 
@@ -22,6 +25,12 @@ class UserRequest(BaseModel):
     k: float = Field(CONFIG["k"])
     layer: int = Field(CONFIG["layer"])
 
+    def model_post_init(self, __context):
+        if self.coeff < 0:
+            self.k *= vector_scaling[self.layer]["k_pos"]
+        else:
+            self.k *= vector_scaling[self.layer]["k_neg"]
+
     def get_api_format(self):
         return {
             "prompt": self.prompt,