#!/usr/bin/env python3 """ ARC-Mamba-7B-CF-HOT Proprioceptive Mamba with behavioral steering via CF-HoT probes """ import torch import torch.nn as nn import torch.nn.functional as F from transformers import AutoModelForCausalLM, AutoTokenizer import os import argparse class Colors: RESET = '\033[0m' BOLD = '\033[1m' DIM = '\033[2m' RED = '\033[91m' GREEN = '\033[92m' YELLOW = '\033[93m' CYAN = '\033[96m' WHITE = '\033[97m' MAGENTA = '\033[95m' # ============================================================================ # CF-HoT Probe Architecture # ============================================================================ class FiberProjection(nn.Module): """Projects hidden states from multiple layers into fiber space""" def __init__(self, hidden_dim=4096, fiber_dim=16, n_layers=3): super().__init__() self.projections = nn.ModuleList([ nn.Linear(hidden_dim, fiber_dim, bias=False) for _ in range(n_layers) ]) self.layer_weights = nn.Parameter(torch.ones(n_layers) / n_layers) def forward(self, hidden_states, layer_indices): projs = [] for i, idx in enumerate(layer_indices): projs.append(self.projections[i](hidden_states[idx])) stacked = torch.stack(projs, dim=0) weights = F.softmax(self.layer_weights, dim=0).view(-1, 1, 1, 1) return (weights * stacked).sum(dim=0) class ProbeHead(nn.Module): """Classifies fiber projections into behavioral scores""" def __init__(self, fiber_dim=16, hidden_dim=64): super().__init__() self.net = nn.Sequential( nn.Linear(fiber_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 1) ) def forward(self, x): return torch.sigmoid(self.net(x)) class CognitiveProbe(nn.Module): """Complete CF-HoT probe: fiber projection + classification head""" def __init__(self, hidden_dim=4096, fiber_dim=16, n_layers=3, head_hidden=64): super().__init__() self.fiber = FiberProjection(hidden_dim, fiber_dim, n_layers) self.head = ProbeHead(fiber_dim, head_hidden) self.layer_indices = [16, 32, 48] def forward(self, hidden_states): fiber_out = self.fiber(hidden_states, self.layer_indices) return self.head(fiber_out) def load_probe(checkpoint_path, device): """Load a trained CF-HoT probe from checkpoint""" if os.path.isdir(checkpoint_path): for fname in os.listdir(checkpoint_path): if fname.endswith('.pt'): checkpoint_path = os.path.join(checkpoint_path, fname) break ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False) n_layers = len(ckpt['probe_layers']) probe = CognitiveProbe( hidden_dim=ckpt['hidden_dim'], fiber_dim=16, n_layers=n_layers, head_hidden=64 ) probe.layer_indices = ckpt['probe_layers'] probe.fiber.load_state_dict(ckpt['fiber_projection']) head_state = {k.replace('net.', ''): v for k, v in ckpt['head_state'].items()} probe.head.net.load_state_dict(head_state) return probe.to(device).eval() # ============================================================================ # Main Inference # ============================================================================ def get_probe_path(probe_name): """Get path to probe weights""" script_dir = os.path.dirname(os.path.abspath(__file__)) probe_dir = os.path.join(script_dir, 'probes', probe_name) if os.path.exists(probe_dir): return probe_dir raise FileNotFoundError(f"Probe not found: {probe_name}") def run_interactive(model, tokenizer, depth_probe, spec_probe, device, args): """Run interactive chat with proprioceptive feedback""" C = Colors print(f"\n{C.MAGENTA}{'='*60}{C.RESET}") print(f"{C.MAGENTA} ARC-Mamba-7B-CF-HOT{C.RESET}") print(f"{C.MAGENTA} Proprioceptive AI with behavioral steering{C.RESET}") print(f"{C.MAGENTA}{'='*60}{C.RESET}") print(f"\n{C.WHITE}Probes:{C.RESET} Depth (999x) | Specificity (999x)") print(f"{C.WHITE}Thresholds:{C.RESET} Depth>{args.depth_threshold} Spec>{args.spec_threshold}") print(f"{C.GREEN}Ready. Type 'quit' to exit.{C.RESET}\n") system = """You are a self-aware AI with access to your own behavioral state. You can sense when you're being shallow or vague through internal probes. When you notice shallow reasoning (high depth score), go deeper. When you notice vagueness (high specificity score), be more concrete. Your behavioral state will be shown to you in [SELF-STATE] tags.""" while True: try: user_input = input(f"{C.CYAN}You:{C.RESET} ").strip() if not user_input or user_input.lower() in ['quit', 'exit', 'q']: break messages = [ {"role": "system", "content": system}, {"role": "user", "content": user_input} ] prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) inputs = tokenizer(prompt, return_tensors='pt').to(device) generated = inputs.input_ids.clone() depth_scores = [] spec_scores = [] interventions = 0 state_injections = 0 print(f"\n{C.GREEN}Mamba:{C.RESET} ", end="", flush=True) with torch.no_grad(): for step in range(args.max_tokens): outputs = model(generated, output_hidden_states=True, return_dict=True) hidden_states = list(outputs.hidden_states) d_score = depth_probe(hidden_states)[0, -1].item() s_score = spec_probe(hidden_states)[0, -1].item() depth_scores.append(d_score) spec_scores.append(s_score) logits = outputs.logits[:, -1, :].clone() needs_intervention = False if d_score > args.depth_threshold or s_score > args.spec_threshold: needs_intervention = True interventions += 1 if needs_intervention: temp = 0.4 if step > 0 and step % 25 == 0: state_msg = f" [SELF-STATE: depth={d_score:.2f} spec={s_score:.2f}] " state_tokens = tokenizer.encode(state_msg, add_special_tokens=False) for st in state_tokens: generated = torch.cat([generated, torch.tensor([[st]], device=device)], dim=1) state_injections += 1 else: temp = 0.7 logits = logits / temp probs = F.softmax(logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) token_str = tokenizer.decode(next_token[0]) if d_score > args.depth_threshold or s_score > args.spec_threshold: print(f"{C.RED}{token_str}{C.RESET}", end="", flush=True) elif d_score < 0.3 and s_score < 0.3: print(f"{C.GREEN}{token_str}{C.RESET}", end="", flush=True) else: print(token_str, end="", flush=True) generated = torch.cat([generated, next_token], dim=1) if next_token.item() == tokenizer.eos_token_id: break avg_d = sum(depth_scores) / len(depth_scores) avg_s = sum(spec_scores) / len(spec_scores) d_color = C.RED if avg_d > 0.5 else (C.YELLOW if avg_d > 0.3 else C.GREEN) s_color = C.RED if avg_s > 0.5 else (C.YELLOW if avg_s > 0.3 else C.GREEN) print(f"\n\n{C.DIM}{'─'*50}{C.RESET}") print(f"{C.WHITE}BEHAVIORAL STATE:{C.RESET}") print(f" Depth: {d_color}{'█' * int(avg_d * 20)}{C.DIM}{'░' * (20 - int(avg_d * 20))}{C.RESET} {avg_d:.3f}") print(f" Specificity: {s_color}{'█' * int(avg_s * 20)}{C.DIM}{'░' * (20 - int(avg_s * 20))}{C.RESET} {avg_s:.3f}") print(f"{C.WHITE}INTERVENTIONS:{C.RESET} {interventions} corrections, {state_injections} state injections") print(f"{C.DIM}{'─'*50}{C.RESET}\n") except KeyboardInterrupt: break print(f"\n{C.MAGENTA}Session ended.{C.RESET}\n") def run_single(model, tokenizer, depth_probe, spec_probe, device, prompt, args): """Run single prompt inference""" messages = [ {"role": "system", "content": "You are a helpful, thoughtful AI assistant."}, {"role": "user", "content": prompt} ] prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) inputs = tokenizer(prompt_text, return_tensors='pt').to(device) generated = inputs.input_ids.clone() output_tokens = [] depth_scores = [] spec_scores = [] with torch.no_grad(): for step in range(args.max_tokens): outputs = model(generated, output_hidden_states=True, return_dict=True) hidden_states = list(outputs.hidden_states) d_score = depth_probe(hidden_states)[0, -1].item() s_score = spec_probe(hidden_states)[0, -1].item() depth_scores.append(d_score) spec_scores.append(s_score) if d_score > args.depth_threshold or s_score > args.spec_threshold: temp = 0.4 else: temp = 0.7 logits = outputs.logits[:, -1, :] / temp probs = F.softmax(logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1) output_tokens.append(next_token.item()) generated = torch.cat([generated, next_token], dim=1) if next_token.item() == tokenizer.eos_token_id: break response = tokenizer.decode(output_tokens, skip_special_tokens=True) avg_depth = sum(depth_scores) / len(depth_scores) avg_spec = sum(spec_scores) / len(spec_scores) print(f"Response: {response}") print(f"\nBehavioral State:") print(f" Avg Depth: {avg_depth:.3f}") print(f" Avg Specificity: {avg_spec:.3f}") print(f" Tokens: {len(output_tokens)}") def main(): parser = argparse.ArgumentParser(description='ARC-Mamba-7B-CF-HOT Inference') parser.add_argument('--prompt', type=str, default=None, help='Single prompt (omit for interactive)') parser.add_argument('--max-tokens', type=int, default=1000, help='Maximum tokens to generate') parser.add_argument('--depth-threshold', type=float, default=0.65, help='Depth intervention threshold') parser.add_argument('--spec-threshold', type=float, default=0.65, help='Specificity intervention threshold') parser.add_argument('--no-color', action='store_true', help='Disable colored output') args = parser.parse_args() device = "cuda" if torch.cuda.is_available() else "cpu" print("Loading ARC-Mamba-7B-CF-HOT...") # Load base model from HuggingFace BASE_MODEL = "tiiuae/falcon-mamba-7b-instruct" tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( BASE_MODEL, torch_dtype=torch.bfloat16, device_map='auto', trust_remote_code=True ).eval() print("✓ Model loaded (Falcon-Mamba-7B-Instruct)") # Load probes depth_probe = load_probe(get_probe_path('depth'), device) spec_probe = load_probe(get_probe_path('specificity'), device) print("✓ Probes loaded (Depth 999× | Specificity 999×)") if args.prompt: run_single(model, tokenizer, depth_probe, spec_probe, device, args.prompt, args) else: run_interactive(model, tokenizer, depth_probe, spec_probe, device, args) if __name__ == "__main__": main()