run.py

#!/usr/bin/env python3
"""
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()