feat: Add playbook vector extraction and activation steering routing to FCES training pipeline
This commit is contained in:
AI-anonymous
473
python/run_representation_pipeline.py
Normal file
473
python/run_representation_pipeline.py
Normal file
@@ -0,0 +1,473 @@
|
||||
"""Orchestration script to compile Skill and Process libraries, train Gating-MLP, and validate steering.
|
||||
|
||||
Discovers local playbooks, extracts representation vectors using Pythia-70m (or dummy fallback),
|
||||
trains the gating network on hidden states, and runs validation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
|
||||
# Import our representation engineering and router modules
|
||||
import sys
|
||||
|
||||
sys.path.append(str(Path(__file__).parent.absolute()))
|
||||
|
||||
from representation_engineering import (
|
||||
PlaybookParser,
|
||||
PlaybookMetadata,
|
||||
RepresentationVectorExtractor,
|
||||
SkillVectorLibrary,
|
||||
ProcessVectorLibrary,
|
||||
)
|
||||
from adapter_moe_router import LearnableGate, ExpertAdapterRouter
|
||||
|
||||
logger = logging.getLogger("run_representation_pipeline")
|
||||
logging.basicConfig(
|
||||
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
|
||||
|
||||
# --- Dummy Model & Tokenizer for Offline Fallback & Testing ---
|
||||
|
||||
|
||||
class DummyTransformerLayer(nn.Module): # type: ignore[misc]
|
||||
def __init__(self, hidden_dim: int) -> None:
|
||||
super().__init__()
|
||||
self.linear1 = nn.Linear(hidden_dim, hidden_dim)
|
||||
self.linear2 = nn.Linear(hidden_dim, hidden_dim)
|
||||
|
||||
def forward(
|
||||
self, x: torch.Tensor, *args: Any, **kwargs: Any
|
||||
) -> Tuple[torch.Tensor]:
|
||||
h = self.linear2(torch.relu(self.linear1(x)))
|
||||
return (x + h,)
|
||||
|
||||
|
||||
class DummyModel(nn.Module): # type: ignore[misc]
|
||||
def __init__(
|
||||
self, vocab_size: int = 1000, hidden_dim: int = 32, num_layers: int = 4
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.vocab_size = vocab_size
|
||||
self.hidden_dim = hidden_dim
|
||||
self.embedding = nn.Embedding(vocab_size, hidden_dim)
|
||||
self.layers = nn.ModuleList(
|
||||
[DummyTransformerLayer(hidden_dim) for _ in range(num_layers)]
|
||||
)
|
||||
self.lm_head = nn.Linear(hidden_dim, vocab_size)
|
||||
|
||||
def forward(self, input_ids: torch.Tensor, *args: Any, **kwargs: Any) -> Any:
|
||||
x = self.embedding(input_ids)
|
||||
for layer in self.layers:
|
||||
x = layer(x)[0]
|
||||
logits = self.lm_head(x)
|
||||
|
||||
class Output:
|
||||
def __init__(self, logits: torch.Tensor) -> None:
|
||||
self.logits = logits
|
||||
|
||||
return Output(logits=logits)
|
||||
|
||||
|
||||
class DummyTokenizer:
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def __call__(
|
||||
self, text: str | List[str], return_tensors: str = "pt", **kwargs: Any
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
if isinstance(text, str):
|
||||
text = [text]
|
||||
batch_ids = []
|
||||
max_len = 0
|
||||
for t in text:
|
||||
words = t.split()
|
||||
ids = [abs(hash(w)) % 1000 for w in words]
|
||||
if not ids:
|
||||
ids = [0]
|
||||
batch_ids.append(ids)
|
||||
max_len = max(max_len, len(ids))
|
||||
|
||||
padded_ids = []
|
||||
for ids in batch_ids:
|
||||
padded_ids.append(ids + [0] * (max_len - len(ids)))
|
||||
return {"input_ids": torch.tensor(padded_ids)}
|
||||
|
||||
|
||||
# --- Helper to parse process steps from playbooks ---
|
||||
|
||||
|
||||
def parse_playbook_steps(path: Path) -> List[Tuple[str, str]]:
|
||||
"""Extracts step names and short descriptions from sequential lists in a playbook."""
|
||||
content = path.read_text(encoding="utf-8")
|
||||
steps = []
|
||||
|
||||
# Locate numbered items under headers like Fallback Procedure, Operational Protocol, etc.
|
||||
lines = content.splitlines()
|
||||
in_procedure = False
|
||||
|
||||
for line in lines:
|
||||
line_str = line.strip()
|
||||
if not line_str:
|
||||
continue
|
||||
|
||||
if line_str.startswith("#"):
|
||||
lower_header = line_str.lower()
|
||||
in_procedure = any(
|
||||
x in lower_header
|
||||
for x in ["procedure", "protocol", "schritt", "ablauf", "loop"]
|
||||
)
|
||||
continue
|
||||
|
||||
if in_procedure:
|
||||
# Match 1. Step name: Description or just 1. Description
|
||||
match = re.match(r"^(\d+)\.\s*(.*?)$", line_str)
|
||||
if match:
|
||||
step_num = match.group(1)
|
||||
step_text = match.group(2).strip()
|
||||
# Split step name and description if separated by double asterisks or colon
|
||||
parts = re.split(r"\*\*|:\s*", step_text, maxsplit=1)
|
||||
if len(parts) > 1 and parts[0].strip():
|
||||
step_name = parts[0].strip()
|
||||
step_desc = parts[1].strip()
|
||||
else:
|
||||
step_name = f"Step_{step_num}"
|
||||
step_desc = step_text
|
||||
steps.append((step_name, step_desc))
|
||||
|
||||
return steps
|
||||
|
||||
|
||||
# --- Pipeline Orchestrator ---
|
||||
|
||||
|
||||
def run_pipeline(
|
||||
model_id: str,
|
||||
skills_dirs: List[str],
|
||||
output_dir: str,
|
||||
use_dummy: bool = False,
|
||||
epochs: int = 50,
|
||||
lr: float = 0.01,
|
||||
) -> None:
|
||||
logger.info("Initializing Representation Engineering Pipeline...")
|
||||
|
||||
# 1. Resolve output directory
|
||||
out_path = Path(output_dir)
|
||||
out_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 2. Load model and tokenizer (with dummy fallback)
|
||||
model = None
|
||||
tokenizer = None
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
if not use_dummy:
|
||||
try:
|
||||
logger.info(f"Attempting to load model '{model_id}' from HuggingFace...")
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
||||
model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
|
||||
logger.info("Model loaded successfully.")
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to load model from HuggingFace: {e}. Falling back to Dummy Model."
|
||||
)
|
||||
use_dummy = True
|
||||
|
||||
if use_dummy:
|
||||
logger.info(
|
||||
"Initializing lightweight Dummy Model and Tokenizer for execution/testing..."
|
||||
)
|
||||
model = DummyModel(vocab_size=1000, hidden_dim=32, num_layers=4).to(device)
|
||||
tokenizer = DummyTokenizer()
|
||||
device = "cpu"
|
||||
|
||||
assert model is not None
|
||||
assert tokenizer is not None
|
||||
|
||||
# 3. Discover playbooks
|
||||
logger.info("Scanning directories for playbooks...")
|
||||
playbooks: List[PlaybookMetadata] = []
|
||||
parsed_files: List[Path] = []
|
||||
|
||||
for s_dir in skills_dirs:
|
||||
path = Path(s_dir)
|
||||
if path.exists():
|
||||
logger.info(f"Scanning directory: {path}")
|
||||
found = PlaybookParser.parse_directory(path)
|
||||
playbooks.extend(found)
|
||||
# Find actual files for process parsing
|
||||
for p_file in path.glob("**/SKILL.md"):
|
||||
parsed_files.append(p_file)
|
||||
for p_file in path.glob("*_SKILL.md"):
|
||||
parsed_files.append(p_file)
|
||||
|
||||
if not playbooks:
|
||||
logger.warning(
|
||||
"No playbooks discovered! Creating a default mock playbook for execution..."
|
||||
)
|
||||
mock_skill = PlaybookMetadata(
|
||||
name="MockSkill",
|
||||
description="A temporary skill for pipeline validation",
|
||||
objectives=[
|
||||
"Validate the routing pipeline",
|
||||
"Verify steering functionality",
|
||||
],
|
||||
trigger_examples=[
|
||||
"Run the mock pipeline check",
|
||||
"Test activation steering on mock",
|
||||
],
|
||||
file_path="mock_SKILL.md",
|
||||
)
|
||||
playbooks.append(mock_skill)
|
||||
|
||||
logger.info(f"Discovered {len(playbooks)} playbooks.")
|
||||
|
||||
# 4. Extract Skill and Process Libraries
|
||||
skill_library = SkillVectorLibrary()
|
||||
process_library = ProcessVectorLibrary()
|
||||
extractor = RepresentationVectorExtractor(model, tokenizer, device=device)
|
||||
|
||||
# Determine model layers
|
||||
transformer_layers = []
|
||||
for name, module in model.named_modules():
|
||||
if re.match(r".*layers?\.\d+$", name):
|
||||
transformer_layers.append(name)
|
||||
num_layers = len(transformer_layers)
|
||||
logger.info(f"Detected {num_layers} transformer layers in base model.")
|
||||
|
||||
# Choose layers to extract: middle/late layers
|
||||
layers_to_extract = (
|
||||
list(range(num_layers // 2, num_layers)) if num_layers > 0 else [0]
|
||||
)
|
||||
|
||||
# Extract skill vectors
|
||||
for pb in playbooks:
|
||||
logger.info(f"Extracting vectors for skill: {pb.name}")
|
||||
vec = extractor.extract_steering_vector(pb, layers_to_extract=layers_to_extract)
|
||||
skill_library.add_vector(vec)
|
||||
|
||||
# Extract sequential process vectors if file exists
|
||||
p_path = Path(pb.file_path)
|
||||
if p_path.exists():
|
||||
steps = parse_playbook_steps(p_path)
|
||||
if steps:
|
||||
logger.info(
|
||||
f"Found {len(steps)} sequential steps in process '{pb.name}'"
|
||||
)
|
||||
step_vectors = []
|
||||
for step_name, step_desc in steps:
|
||||
step_pb = PlaybookMetadata(
|
||||
name=f"{pb.name}_{step_name}",
|
||||
description=step_desc,
|
||||
objectives=[step_desc],
|
||||
trigger_examples=pb.trigger_examples,
|
||||
file_path=str(p_path),
|
||||
)
|
||||
step_vec = extractor.extract_steering_vector(
|
||||
step_pb, layers_to_extract=layers_to_extract
|
||||
)
|
||||
step_vectors.append(step_vec)
|
||||
process_library.add_process(
|
||||
pb.name.lower().replace(" ", "_"), step_vectors
|
||||
)
|
||||
|
||||
# Save libraries
|
||||
skill_lib_path = out_path / "skill_library.pt"
|
||||
process_lib_path = out_path / "process_library.pt"
|
||||
skill_library.save(skill_lib_path)
|
||||
process_library.save(process_lib_path)
|
||||
|
||||
# 5. Train Gating-MLP on hidden states
|
||||
skill_ids = sorted(list(skill_library.vectors.keys()))
|
||||
if not skill_ids:
|
||||
logger.error("No skills available to train Gating-MLP.")
|
||||
return
|
||||
|
||||
logger.info(f"Training Gating-MLP router over {len(skill_ids)} skills...")
|
||||
|
||||
# Gating features layer (we extract hidden states from early/middle layer to predict routing)
|
||||
gate_layer_idx = layers_to_extract[0] if layers_to_extract else 0
|
||||
hidden_dim = model.hidden_dim if hasattr(model, "hidden_dim") else 32
|
||||
if hasattr(model, "config") and hasattr(model.config, "hidden_size"):
|
||||
hidden_dim = model.config.hidden_size
|
||||
|
||||
gate_net = LearnableGate(in_features=hidden_dim, num_adapters=len(skill_ids)).to(
|
||||
device
|
||||
)
|
||||
optimizer = optim.Adam(gate_net.parameters(), lr=lr)
|
||||
criterion = nn.CrossEntropyLoss()
|
||||
|
||||
# Collect training dataset: (hidden_state, label)
|
||||
# We pass the win prompts for each skill, hook the gate layer, and collect states
|
||||
training_data: List[Tuple[torch.Tensor, int]] = []
|
||||
|
||||
# Temporary hook to collect states
|
||||
collected_states: List[torch.Tensor] = []
|
||||
|
||||
def collect_hook(module: nn.Module, input_t: Any, output_t: Any) -> None:
|
||||
x = output_t[0] if isinstance(output_t, tuple) else output_t
|
||||
# Pool to sequence mean
|
||||
collected_states.append(x.detach().mean(dim=1).squeeze(0))
|
||||
|
||||
# Register hook on gate layer
|
||||
hook_handle = None
|
||||
target_layer_name = (
|
||||
transformer_layers[gate_layer_idx]
|
||||
if gate_layer_idx < len(transformer_layers)
|
||||
else None
|
||||
)
|
||||
|
||||
if target_layer_name:
|
||||
for name, module in model.named_modules():
|
||||
if name == target_layer_name:
|
||||
hook_handle = module.register_forward_hook(collect_hook)
|
||||
break
|
||||
|
||||
# Run win prompts to collect activations
|
||||
for label_idx, skill_id in enumerate(skill_ids):
|
||||
# Retrieve parsed metadata for this skill ID
|
||||
pb_match = next(
|
||||
(p for p in playbooks if p.name.lower().replace(" ", "_") == skill_id), None
|
||||
)
|
||||
if pb_match:
|
||||
for trigger in pb_match.trigger_examples:
|
||||
win_prompt = (
|
||||
f"Instructions: You are acting with the following skill: {pb_match.name}.\n"
|
||||
f"Request: {trigger}\nOutput:"
|
||||
)
|
||||
inputs = tokenizer(win_prompt, return_tensors="pt")
|
||||
# Move inputs to device
|
||||
inputs = {k: v.to(device) for k, v in inputs.items()}
|
||||
|
||||
collected_states.clear()
|
||||
with torch.no_grad():
|
||||
model(**inputs)
|
||||
|
||||
if collected_states:
|
||||
state = collected_states[0].cpu() # shape [hidden_dim]
|
||||
training_data.append((state, label_idx))
|
||||
|
||||
if hook_handle:
|
||||
hook_handle.remove()
|
||||
|
||||
if not training_data:
|
||||
logger.warning(
|
||||
"Could not collect training data. Using synthetic data to train Gating-MLP."
|
||||
)
|
||||
# Fallback to random features for testing compilation flow
|
||||
for i in range(100):
|
||||
label = i % len(skill_ids)
|
||||
feat = torch.randn(hidden_dim) + (label * 2.0) # separate them a bit
|
||||
training_data.append((feat, label))
|
||||
|
||||
# Train MLP
|
||||
gate_net.train()
|
||||
X = torch.stack([x for x, y in training_data]).to(device)
|
||||
Y = torch.tensor([y for x, y in training_data]).to(device)
|
||||
|
||||
logger.info(
|
||||
f"Collected {len(training_data)} training samples. Starting optimization..."
|
||||
)
|
||||
|
||||
for epoch in range(epochs):
|
||||
optimizer.zero_grad()
|
||||
outputs = gate_net(X)
|
||||
loss = criterion(outputs, Y)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
if (epoch + 1) % max(1, epochs // 5) == 0 or epoch == epochs - 1:
|
||||
# Calculate accuracy
|
||||
_, preds = torch.max(outputs, 1)
|
||||
correct = (preds == Y).sum().item()
|
||||
acc = correct / len(Y)
|
||||
logger.info(
|
||||
f"Epoch {epoch+1:02d}/{epochs:02d} | Loss: {loss.item():.4f} | Training Accuracy: {acc * 100:.1f}%"
|
||||
)
|
||||
|
||||
# Save gate weights
|
||||
gate_weights_path = out_path / "gate_weights.pt"
|
||||
torch.save(gate_net.state_dict(), gate_weights_path)
|
||||
logger.info(f"Gating MLP weights saved to {gate_weights_path}")
|
||||
|
||||
# 6. Validate steerability & routing performance
|
||||
logger.info("Running pipeline verification...")
|
||||
router = ExpertAdapterRouter(
|
||||
base_model=model,
|
||||
skill_library=skill_library,
|
||||
process_library=process_library,
|
||||
in_features=hidden_dim,
|
||||
steering_alpha=1.0,
|
||||
)
|
||||
router.skill_gate.load_state_dict(
|
||||
torch.load(gate_weights_path, map_location=device)
|
||||
)
|
||||
router.register_hooks()
|
||||
|
||||
# Run test prompt with routing enabled
|
||||
test_prompt = "Validate this test prompt"
|
||||
inputs = tokenizer(test_prompt, return_tensors="pt")
|
||||
inputs = {k: v.to(device) for k, v in inputs.items()}
|
||||
|
||||
# Test that forwarding goes through hooks without crash
|
||||
with torch.no_grad():
|
||||
_ = model(**inputs)
|
||||
logger.info("Successfully executed forward pass with dynamic activation steering.")
|
||||
|
||||
router.unregister_hooks()
|
||||
logger.info("Pipeline run completed successfully.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="FCES Representation Engineering Pipeline"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_id",
|
||||
type=str,
|
||||
default="EleutherAI/pythia-70m",
|
||||
help="HuggingFace model identifier",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skills_dirs",
|
||||
type=str,
|
||||
default="C:/Users/Sven/Documents/svenco-knowledge/skills,C:/Users/Sven/Documents/everything-claude-code/skills",
|
||||
help="Comma-separated paths to search for playbooks",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_dir",
|
||||
type=str,
|
||||
default="./python/output",
|
||||
help="Directory to save compiled libraries",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_dummy", action="store_true", help="Force using lightweight dummy model"
|
||||
)
|
||||
parser.add_argument("--epochs", type=int, default=50, help="Gating training epochs")
|
||||
parser.add_argument(
|
||||
"--lr", type=float, default=0.01, help="Gating training learning rate"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Split skills dirs
|
||||
dirs_list = [d.strip() for d in args.skills_dirs.split(",") if d.strip()]
|
||||
|
||||
run_pipeline(
|
||||
model_id=args.model_id,
|
||||
skills_dirs=dirs_list,
|
||||
output_dir=args.output_dir,
|
||||
use_dummy=args.use_dummy,
|
||||
epochs=args.epochs,
|
||||
lr=args.lr,
|
||||
)
|
||||
Reference in New Issue
Block a user