3622b896a0
Three readability fixes for the F8 screen: * Z-score values per-layer by default (`[z]` toggles to raw dot- product). Raw values are dominated by residual-stream magnitude — z-scores read as "σ above concept-vector baseline" which is interpretable and scale-stable across frames. * Stable ordering with TOP_K + HYSTERESIS hysteresis band. Pinned concept set only rotates when a member drops out of the hysteresis band by |value| rank — bars update values in place without names flickering row-to-row. * Default to the deepest hooked layer (index 3 = layer 58 of 64). Clustering validation showed layer 58 is the only one with strong within-family cohesion (fear +0.37, shame +0.29, sadness +0.25 cosine); earlier layers are mostly noise for this task. Co-Authored-By: Proof of Concept <poc@bcachefs.org>
399 lines
15 KiB
Rust
399 lines
15 KiB
Rust
// amygdala.rs — F8 amygdala screen: live per-token concept-readout
|
||
// projections from the vLLM server's readout.safetensors.
|
||
//
|
||
// Left panel: top-K concepts by magnitude at the currently-selected
|
||
// layer, as horizontal bars. The concept names come from the manifest
|
||
// fetched at agent startup; the values come from the per-token readout
|
||
// pushed onto agent.readout by the streaming token handler.
|
||
//
|
||
// Bottom: scrolling history of the last few tokens' top concept.
|
||
//
|
||
// Keys:
|
||
// 1..9 select layer index (1 = first layer in the manifest)
|
||
// t toggle between "current" (last token) and "mean over recent"
|
||
|
||
use ratatui::{
|
||
layout::{Constraint, Direction, Layout, Rect},
|
||
style::{Color, Modifier, Style},
|
||
text::{Line, Span},
|
||
widgets::{Block, Borders, Gauge, Paragraph, Wrap},
|
||
Frame,
|
||
};
|
||
use ratatui::crossterm::event::{Event, KeyCode};
|
||
|
||
use super::{App, ScreenView};
|
||
use crate::agent::api::ReadoutManifest;
|
||
use crate::agent::readout::ReadoutEntry;
|
||
|
||
const TOP_K: usize = 20;
|
||
/// Hysteresis band around TOP_K. A concept currently in the display
|
||
/// is kept as long as its |z-score| rank stays in the top
|
||
/// ``TOP_K + HYSTERESIS``; only falls out when it drops below that.
|
||
/// Prevents the ticker-tape flicker that pure top-K sorting produces.
|
||
const HYSTERESIS: usize = 20;
|
||
|
||
pub(crate) struct AmygdalaScreen {
|
||
selected_layer: usize,
|
||
mode: DisplayMode,
|
||
/// Concept indices currently pinned in display order. Values at
|
||
/// these indices change every frame; the set only rotates when a
|
||
/// pinned concept drops out of the hysteresis band.
|
||
display_indices: Vec<usize>,
|
||
/// Whether to show z-scored values (default) or raw dot products.
|
||
normalize: bool,
|
||
}
|
||
|
||
#[derive(Clone, Copy, PartialEq)]
|
||
enum DisplayMode {
|
||
/// Values from the single most recent token.
|
||
Current,
|
||
/// Mean over all tokens currently in the ring buffer.
|
||
MeanRecent,
|
||
}
|
||
|
||
impl AmygdalaScreen {
|
||
pub fn new() -> Self {
|
||
Self {
|
||
// Default to the deepest hooked layer — emotion/concept
|
||
// circuits concentrate in the last ~20% of the network,
|
||
// and our clustering validation showed layer 58 was the
|
||
// only one with strong within-family cohesion. Bounded
|
||
// down to the actual layer count at render time.
|
||
selected_layer: 3,
|
||
mode: DisplayMode::MeanRecent,
|
||
display_indices: Vec::new(),
|
||
normalize: true,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl ScreenView for AmygdalaScreen {
|
||
fn label(&self) -> &'static str { "amygdala" }
|
||
|
||
fn tick(&mut self, frame: &mut Frame, area: Rect,
|
||
events: &[Event], app: &mut App) {
|
||
for event in events {
|
||
if let Event::Key(key) = event {
|
||
match key.code {
|
||
KeyCode::Char(c) if c.is_ascii_digit() && c != '0' => {
|
||
let idx = (c as u8 - b'1') as usize;
|
||
self.selected_layer = idx;
|
||
}
|
||
KeyCode::Char('t') => {
|
||
self.mode = match self.mode {
|
||
DisplayMode::Current => DisplayMode::MeanRecent,
|
||
DisplayMode::MeanRecent => DisplayMode::Current,
|
||
};
|
||
// Re-pin on mode change; the relative
|
||
// magnitudes between current-token and
|
||
// mean-recent differ substantially.
|
||
self.display_indices.clear();
|
||
}
|
||
KeyCode::Char('z') => {
|
||
self.normalize = !self.normalize;
|
||
self.display_indices.clear();
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Snapshot the shared buffer with a short lock.
|
||
let snapshot = match app.agent.readout.lock() {
|
||
Ok(buf) => {
|
||
if !buf.is_enabled() {
|
||
render_disabled(frame, area);
|
||
return;
|
||
}
|
||
let manifest = buf.manifest.clone().unwrap();
|
||
let entries: Vec<ReadoutEntry> =
|
||
buf.recent.iter().cloned().collect();
|
||
(manifest, entries)
|
||
}
|
||
Err(_) => {
|
||
render_disabled(frame, area);
|
||
return;
|
||
}
|
||
};
|
||
let (manifest, entries) = snapshot;
|
||
|
||
// Bound the selected layer to what the manifest actually has.
|
||
let n_layers = manifest.layers.len();
|
||
if self.selected_layer >= n_layers {
|
||
self.selected_layer = 0;
|
||
}
|
||
|
||
// Compute the raw values for the selected layer: either the
|
||
// latest token's row, or the mean across recent tokens. Raw
|
||
// means un-normalized dot products — their absolute scale is
|
||
// dominated by residual-stream norm, not concept alignment.
|
||
let raw: Option<Vec<f32>> = match self.mode {
|
||
DisplayMode::Current => entries
|
||
.last()
|
||
.and_then(|e| e.readout.get(self.selected_layer).cloned()),
|
||
DisplayMode::MeanRecent => mean_layer(&entries, self.selected_layer),
|
||
};
|
||
|
||
// Optional z-score normalization: remove the per-layer mean,
|
||
// scale by std. Result is "σ above/below the concept-vector
|
||
// average at this layer" — the loud-residual-stream scaling
|
||
// factor cancels out, values become comparable across frames.
|
||
let display_values = raw.as_ref().map(|v| {
|
||
if self.normalize { z_score(v) } else { v.clone() }
|
||
});
|
||
|
||
// Update the pinned display set with hysteresis: a concept
|
||
// stays pinned while it remains in the top (TOP_K + HYSTERESIS)
|
||
// by |value|; falls out only when it drops below that band.
|
||
// Keeps rows stable while values update in place.
|
||
if let Some(v) = display_values.as_ref() {
|
||
self.refresh_display_indices(v);
|
||
}
|
||
|
||
let layout = Layout::default()
|
||
.direction(Direction::Vertical)
|
||
.constraints([
|
||
Constraint::Length(3), // header
|
||
Constraint::Min(10), // bars
|
||
Constraint::Length(6), // recent tokens
|
||
])
|
||
.split(area);
|
||
|
||
render_header(frame, layout[0], &manifest, self.selected_layer,
|
||
self.mode, entries.len(), self.normalize);
|
||
match display_values {
|
||
Some(v) => render_bars(
|
||
frame, layout[1], &manifest.concepts, &v,
|
||
&self.display_indices, self.normalize,
|
||
),
|
||
None => render_empty_bars(frame, layout[1]),
|
||
}
|
||
render_recent(frame, layout[2], &entries, self.selected_layer,
|
||
&manifest.concepts);
|
||
}
|
||
}
|
||
|
||
impl AmygdalaScreen {
|
||
/// Add concepts entering the hysteresis band; evict concepts that
|
||
/// dropped out. Preserves existing order for concepts that stay.
|
||
fn refresh_display_indices(&mut self, values: &[f32]) {
|
||
let n = values.len();
|
||
if n == 0 {
|
||
return;
|
||
}
|
||
// Rank all concepts by |value| desc so we can check both "in
|
||
// strict top-K" and "in hysteresis band (top K + H)" cheaply.
|
||
let mut rank: Vec<(usize, f32)> = values.iter()
|
||
.enumerate().map(|(i, v)| (i, v.abs())).collect();
|
||
rank.sort_by(|a, b| b.1.partial_cmp(&a.1)
|
||
.unwrap_or(std::cmp::Ordering::Equal));
|
||
let hyst_cutoff = (TOP_K + HYSTERESIS).min(n);
|
||
let in_band: std::collections::HashSet<usize> =
|
||
rank.iter().take(hyst_cutoff).map(|(i, _)| *i).collect();
|
||
// Drop anything that left the band.
|
||
self.display_indices.retain(|i| in_band.contains(i));
|
||
// Fill up to TOP_K by walking the top-K-by-|value| and adding
|
||
// any concept not already displayed.
|
||
for (i, _) in rank.iter().take(TOP_K) {
|
||
if self.display_indices.len() >= TOP_K {
|
||
break;
|
||
}
|
||
if !self.display_indices.contains(i) {
|
||
self.display_indices.push(*i);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
fn render_disabled(frame: &mut Frame, area: Rect) {
|
||
let text = Paragraph::new(Line::from(vec![
|
||
Span::raw("readout disabled — server did not return a manifest. "),
|
||
Span::styled("Start vLLM with ", Style::default().fg(Color::DarkGray)),
|
||
Span::styled("VLLM_READOUT_MANIFEST", Style::default().fg(Color::Yellow)),
|
||
Span::styled(" + ", Style::default().fg(Color::DarkGray)),
|
||
Span::styled("VLLM_READOUT_VECTORS", Style::default().fg(Color::Yellow)),
|
||
Span::styled(".", Style::default().fg(Color::DarkGray)),
|
||
]))
|
||
.wrap(Wrap { trim: true })
|
||
.block(Block::default().borders(Borders::ALL).title("amygdala"));
|
||
frame.render_widget(text, area);
|
||
}
|
||
|
||
fn render_header(frame: &mut Frame, area: Rect, manifest: &ReadoutManifest,
|
||
selected: usize, mode: DisplayMode, n_tokens: usize,
|
||
normalize: bool) {
|
||
let mode_str = match mode {
|
||
DisplayMode::Current => "current",
|
||
DisplayMode::MeanRecent => "mean(recent)",
|
||
};
|
||
let scale_str = if normalize { "z-score" } else { "raw" };
|
||
let layer = manifest.layers.get(selected).copied().unwrap_or(0);
|
||
let spans = vec![
|
||
Span::styled("layer ", Style::default().fg(Color::DarkGray)),
|
||
Span::styled(
|
||
format!("{}/{} ", selected + 1, manifest.layers.len()),
|
||
Style::default().add_modifier(Modifier::BOLD),
|
||
),
|
||
Span::styled("(index ", Style::default().fg(Color::DarkGray)),
|
||
Span::styled(format!("{}", layer), Style::default().fg(Color::Cyan)),
|
||
Span::styled(") ", Style::default().fg(Color::DarkGray)),
|
||
Span::styled("mode ", Style::default().fg(Color::DarkGray)),
|
||
Span::styled(mode_str, Style::default().fg(Color::Yellow)),
|
||
Span::styled(" scale ", Style::default().fg(Color::DarkGray)),
|
||
Span::styled(scale_str, Style::default().fg(Color::Yellow)),
|
||
Span::styled(" ", Style::default()),
|
||
Span::styled(
|
||
format!("{} toks in ring", n_tokens),
|
||
Style::default().fg(Color::DarkGray),
|
||
),
|
||
Span::raw(" "),
|
||
Span::styled(
|
||
format!("[1-{}] layer [t] mode [z] z-score/raw",
|
||
manifest.layers.len().min(9)),
|
||
Style::default().fg(Color::DarkGray),
|
||
),
|
||
];
|
||
let para = Paragraph::new(Line::from(spans))
|
||
.block(Block::default().borders(Borders::ALL).title("amygdala"));
|
||
frame.render_widget(para, area);
|
||
}
|
||
|
||
fn render_bars(frame: &mut Frame, area: Rect,
|
||
concepts: &[String], values: &[f32],
|
||
display_indices: &[usize], normalize: bool) {
|
||
let inner = Block::default().borders(Borders::ALL)
|
||
.title("top concepts");
|
||
let inner_area = inner.inner(area);
|
||
frame.render_widget(inner, area);
|
||
|
||
if inner_area.height == 0 || display_indices.is_empty() {
|
||
return;
|
||
}
|
||
|
||
// Bar-scale normalization. For z-score mode, pin the bar to a
|
||
// fixed reference (|z| = 3 = full bar) so the visual magnitude
|
||
// has a meaningful interpretation ("3σ from baseline"). For raw
|
||
// mode, fall back to the old behavior (scale to the loudest
|
||
// concept on-screen).
|
||
let scale_ref: f32 = if normalize {
|
||
3.0
|
||
} else {
|
||
display_indices.iter()
|
||
.filter_map(|&i| values.get(i))
|
||
.map(|v| v.abs())
|
||
.fold(0.0_f32, f32::max)
|
||
.max(1e-6)
|
||
};
|
||
|
||
let rows = (inner_area.height as usize).min(display_indices.len());
|
||
let row_constraints: Vec<Constraint> =
|
||
std::iter::repeat(Constraint::Length(1)).take(rows).collect();
|
||
let chunks = Layout::default()
|
||
.direction(Direction::Vertical)
|
||
.constraints(row_constraints)
|
||
.split(inner_area);
|
||
|
||
for (row, &c_idx) in display_indices.iter().take(rows).enumerate() {
|
||
let v = values.get(c_idx).copied().unwrap_or(0.0);
|
||
let label = concepts.get(c_idx).cloned()
|
||
.unwrap_or_else(|| format!("c{}", c_idx));
|
||
let ratio = (v.abs() / scale_ref).clamp(0.0, 1.0);
|
||
let color = if v >= 0.0 { Color::Green } else { Color::Red };
|
||
let display_num = if normalize {
|
||
format!("{:+.2}σ", v)
|
||
} else {
|
||
format!("{:+.3}", v)
|
||
};
|
||
let gauge = Gauge::default()
|
||
.ratio(ratio as f64)
|
||
.gauge_style(Style::default().fg(color).bg(Color::Reset))
|
||
.label(format!("{:<26} {}", truncate_name(&label, 26), display_num));
|
||
frame.render_widget(gauge, chunks[row]);
|
||
}
|
||
}
|
||
|
||
fn render_empty_bars(frame: &mut Frame, area: Rect) {
|
||
let para = Paragraph::new(Line::from(Span::styled(
|
||
"waiting for tokens…",
|
||
Style::default().fg(Color::DarkGray),
|
||
)))
|
||
.block(Block::default().borders(Borders::ALL).title("top concepts"));
|
||
frame.render_widget(para, area);
|
||
}
|
||
|
||
fn render_recent(frame: &mut Frame, area: Rect, entries: &[ReadoutEntry],
|
||
layer: usize, concepts: &[String]) {
|
||
let mut lines: Vec<Line> = Vec::new();
|
||
for entry in entries.iter().rev().take(4) {
|
||
let row = match entry.readout.get(layer) {
|
||
Some(r) => r,
|
||
None => continue,
|
||
};
|
||
// top concept at this layer for this token
|
||
let (best_idx, best_val) = row.iter().enumerate()
|
||
.fold((0, 0.0_f32), |acc, (i, v)| {
|
||
if v.abs() > acc.1.abs() { (i, *v) } else { acc }
|
||
});
|
||
let name = concepts.get(best_idx).cloned()
|
||
.unwrap_or_else(|| format!("c{}", best_idx));
|
||
let tok_str = format!("t{:>5}", entry.token_id);
|
||
lines.push(Line::from(vec![
|
||
Span::styled(tok_str, Style::default().fg(Color::DarkGray)),
|
||
Span::raw(" "),
|
||
Span::styled(
|
||
format!("{:<24}", truncate_name(&name, 24)),
|
||
Style::default().fg(
|
||
if best_val >= 0.0 { Color::Green } else { Color::Red },
|
||
),
|
||
),
|
||
Span::styled(
|
||
format!(" {:+.3}", best_val),
|
||
Style::default().add_modifier(Modifier::BOLD),
|
||
),
|
||
]));
|
||
}
|
||
let para = Paragraph::new(lines)
|
||
.block(Block::default().borders(Borders::ALL).title("recent tokens — top concept"));
|
||
frame.render_widget(para, area);
|
||
}
|
||
|
||
/// Z-score normalize: `(v - mean) / std` across the concept axis.
|
||
/// Result is comparable across frames and layers (the residual-stream
|
||
/// magnitude factors out) and has the nice property that "this is
|
||
/// ≥2σ elevated" has a concrete meaning regardless of scale.
|
||
fn z_score(values: &[f32]) -> Vec<f32> {
|
||
let n = values.len() as f32;
|
||
if n == 0.0 {
|
||
return Vec::new();
|
||
}
|
||
let mean = values.iter().sum::<f32>() / n;
|
||
let var = values.iter()
|
||
.map(|v| (v - mean) * (v - mean))
|
||
.sum::<f32>() / n;
|
||
let std = var.sqrt().max(1e-6);
|
||
values.iter().map(|v| (v - mean) / std).collect()
|
||
}
|
||
|
||
fn mean_layer(entries: &[ReadoutEntry], layer: usize) -> Option<Vec<f32>> {
|
||
let rows: Vec<&Vec<f32>> = entries.iter()
|
||
.filter_map(|e| e.readout.get(layer))
|
||
.collect();
|
||
if rows.is_empty() {
|
||
return None;
|
||
}
|
||
let n_concepts = rows[0].len();
|
||
let mut acc = vec![0.0_f32; n_concepts];
|
||
for r in &rows {
|
||
for (i, v) in r.iter().enumerate() {
|
||
acc[i] += *v;
|
||
}
|
||
}
|
||
let n = rows.len() as f32;
|
||
for v in &mut acc { *v /= n; }
|
||
Some(acc)
|
||
}
|
||
|
||
fn truncate_name(s: &str, max: usize) -> String {
|
||
if s.len() <= max { s.to_string() }
|
||
else { format!("{}…", &s[..max.saturating_sub(1)]) }
|
||
}
|