frankaging commited on
Commit
906c92b
·
1 Parent(s): 5d1859a
Files changed (1) hide show
  1. app.py +13 -15
app.py CHANGED
@@ -14,7 +14,7 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
14
  login(token=HF_TOKEN)
15
 
16
  MAX_MAX_NEW_TOKENS = 2048
17
- DEFAULT_MAX_NEW_TOKENS = 256 # smaller default to save memory
18
  MAX_INPUT_TOKEN_LENGTH = 4096
19
 
20
  css = """
@@ -178,7 +178,6 @@ def generate(
178
  }
179
  ] if subspaces_list else None,
180
  "streamer": streamer,
181
- "repetition_penalty": 1.5,
182
  "do_sample": True
183
  }
184
 
@@ -222,8 +221,8 @@ def update_dropdown_choices(search_text):
222
  filtered = filter_concepts(search_text)
223
  if not filtered or len(filtered) == 0:
224
  return gr.update(choices=[f"[New] {search_text}"], value=f"[New] {search_text}", interactive=True), gr.Textbox(
225
- label="No matching existing concepts were found!",
226
- value="Good news! Based on the concept you provided, we will automatically generate a steering vector. Try it out by starting a chat!",
227
  lines=3,
228
  interactive=False,
229
  visible=True,
@@ -240,43 +239,42 @@ with gr.Blocks(css=css, fill_height=True) as demo:
240
  # Remove default subspaces
241
  selected_subspaces = gr.State([])
242
 
243
- with gr.Row(min_height=700):
244
  # Left side: bigger chat area
245
  with gr.Column(scale=7):
246
  chat_interface = gr.ChatInterface(
247
  fn=generate,
248
- title="Chat with a Concept Steering Model",
249
- description="""Steer responses by selecting concepts on the right →\n\nWe intervene on Gemma-2-2B-it by adding steering vectors to the residual stream at layer 20.""",
250
  type="messages",
251
  additional_inputs=[selected_subspaces],
252
- fill_height=True
253
  )
254
 
255
  # Right side: concept management
256
  with gr.Column(scale=3):
257
- gr.Markdown("## Steer Model Responses")
258
- gr.Markdown("Search and then select a concept to steer. The closest match will be automatically selected. If there is no match, our steering vector generator (a finetuned Gemma-2-2B) auto-steers for you!")
259
  # Concept Search and Selection
260
  with gr.Group():
261
  search_box = gr.Textbox(
262
- label="Search Concepts",
263
- placeholder="Find concepts to steer the model (e.g. 'time travel')",
264
  lines=2,
265
  )
266
  msg = gr.TextArea(visible=False)
267
  concept_dropdown = gr.Dropdown(
268
- label="Select a concept to steer the model (Click to see more!)",
269
  interactive=True,
270
  allow_custom_value=False,
271
  )
272
  concept_magnitude = gr.Slider(
273
- label="Steering Intensity",
274
  minimum=-5,
275
  maximum=5,
276
  step=0.1,
277
  value=3,
278
  )
279
-
280
  # Wire up events
281
  # When search box changes, update dropdown AND trigger concept selection
282
  search_box.input(
 
14
  login(token=HF_TOKEN)
15
 
16
  MAX_MAX_NEW_TOKENS = 2048
17
+ DEFAULT_MAX_NEW_TOKENS = 128 # smaller default to save memory
18
  MAX_INPUT_TOKEN_LENGTH = 4096
19
 
20
  css = """
 
178
  }
179
  ] if subspaces_list else None,
180
  "streamer": streamer,
 
181
  "do_sample": True
182
  }
183
 
 
221
  filtered = filter_concepts(search_text)
222
  if not filtered or len(filtered) == 0:
223
  return gr.update(choices=[f"[New] {search_text}"], value=f"[New] {search_text}", interactive=True), gr.Textbox(
224
+ label="No matching existing topics were found!",
225
+ value="Good news! Based on the topic you provided, we will automatically generate a steering vector. Try it out by starting a chat!",
226
  lines=3,
227
  interactive=False,
228
  visible=True,
 
239
  # Remove default subspaces
240
  selected_subspaces = gr.State([])
241
 
242
+ with gr.Row(min_height=300):
243
  # Left side: bigger chat area
244
  with gr.Column(scale=7):
245
  chat_interface = gr.ChatInterface(
246
  fn=generate,
247
+ title="Chat with a Topic Steering Model",
248
+ description="""Choose a topic you want the model to discuss on the right →\n\nWe intervene on Gemma-2-2B-it by adding steering vectors to the residual stream at layer 20.""",
249
  type="messages",
250
  additional_inputs=[selected_subspaces],
 
251
  )
252
 
253
  # Right side: concept management
254
  with gr.Column(scale=3):
255
+ gr.Markdown("# Steer model responses")
256
+ gr.Markdown("Search and then select a topic you want the model to discuss. The closest match will be automatically selected. If there is no match, a finetuned Gemma-2-2B model auto-steers for you!")
257
  # Concept Search and Selection
258
  with gr.Group():
259
  search_box = gr.Textbox(
260
+ label="Search topics to steer",
261
+ placeholder="Try: 'time travel'",
262
  lines=2,
263
  )
264
  msg = gr.TextArea(visible=False)
265
  concept_dropdown = gr.Dropdown(
266
+ label="Select a topic to steer the model (Click to see more!)",
267
  interactive=True,
268
  allow_custom_value=False,
269
  )
270
  concept_magnitude = gr.Slider(
271
+ label="Steering intensity",
272
  minimum=-5,
273
  maximum=5,
274
  step=0.1,
275
  value=3,
276
  )
277
+
278
  # Wire up events
279
  # When search box changes, update dropdown AND trigger concept selection
280
  search_box.input(