| |
| """ |
| MAMBA CHAT WITH SELF-AWARE CF-HoT INTERVENTION |
| The model reads its own behavioral state and steers itself |
| """ |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| import os |
|
|
| class C: |
| RESET = '\033[0m' |
| BOLD = '\033[1m' |
| DIM = '\033[2m' |
| RED = '\033[91m' |
| GREEN = '\033[92m' |
| YELLOW = '\033[93m' |
| CYAN = '\033[96m' |
| WHITE = '\033[97m' |
|
|
| class FiberProjection(nn.Module): |
| 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): |
| 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): |
| 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): |
| 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() |
|
|
| def main(): |
| print(f"\n{C.CYAN}{'═'*60}{C.RESET}") |
| print(f"{C.CYAN} PROPRIOCEPTIVE MAMBA - SELF-AWARE INFERENCE{C.RESET}") |
| print(f"{C.CYAN} Model reads its own behavioral state and self-corrects{C.RESET}") |
| print(f"{C.CYAN}{'═'*60}{C.RESET}\n") |
| |
| device = "cuda" |
| |
| DEPTH_THRESHOLD = 0.65 |
| SPEC_THRESHOLD = 0.65 |
| |
| print(f"{C.WHITE}Loading Falcon-Mamba-7B...{C.RESET}") |
| tokenizer = AutoTokenizer.from_pretrained('tiiuae/falcon-mamba-7b-instruct', trust_remote_code=True) |
| model = AutoModelForCausalLM.from_pretrained( |
| 'tiiuae/falcon-mamba-7b-instruct', |
| torch_dtype=torch.bfloat16, |
| device_map='auto', |
| trust_remote_code=True |
| ) |
| model.eval() |
| print(f"{C.GREEN}✓ Model loaded{C.RESET}") |
| |
| print(f"{C.WHITE}Loading behavioral probes...{C.RESET}") |
| depth_probe = load_probe("/home/programmer/Desktop/Claude_and_me/mamba7b_cognitive_output/depth/ckpt_1000", device) |
| spec_probe = load_probe("/home/programmer/Desktop/Claude_and_me/mamba7b_cognitive_output/specificity/ckpt_1000", device) |
| print(f"{C.GREEN}✓ Probes loaded (Depth 999× | Specificity 999×){C.RESET}") |
| |
| print(f"\n{C.YELLOW}Intervention thresholds: Depth>{DEPTH_THRESHOLD} Spec>{SPEC_THRESHOLD}{C.RESET}") |
| 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(200): |
| 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 > DEPTH_THRESHOLD or s_score > 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 > DEPTH_THRESHOLD or s_score > 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.CYAN}Proprioceptive AI session complete.{C.RESET}\n") |
|
|
| if __name__ == "__main__": |
| main() |
