|
|
|
|
|
""" |
|
|
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() |
