File size: 8,870 Bytes
c2d8a57 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 | #!/usr/bin/env python3
"""
PC-SHO-DLM Chat — Conditional masked diffusion chatbot.
The prompt stays unmasked. The response region starts fully masked.
PC settling iteratively denoises the response conditioned on the prompt.
Usage:
python3 train_chat.py --data alpaca # Train on dialogue data
python3 chat.py --checkpoint checkpoints/chat/chat_final.pt # Chat
"""
import argparse, glob, os, sys, torch
import torch.nn.functional as F
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))
from model import PCSHODLM, PCSHOConfig, LocalParameterUpdater
def find_checkpoint():
for p in ["checkpoints/chat/chat_final.pt", "checkpoints/chat/*.pt",
"checkpoints/*/checkpoint_final.pt", "checkpoints/*/*.pt"]:
files = sorted(glob.glob(p), key=os.path.getmtime, reverse=True)
if files:
return files[0]
return None
def encode_text(text, vocab_size=257):
return [min(b + 1, vocab_size - 1) for b in text.encode("utf-8")]
def decode_tokens(tokens):
raw = tokens if isinstance(tokens, list) else tokens.squeeze().cpu().tolist()
# Filter out mask tokens (0) and padding
return bytes([max(0, t - 1) for t in raw if t > 0 and t < 257]).decode("utf-8", errors="replace")
def generate_response(model, prompt_text, response_len=256, device="cpu",
temperature=0.8, n_steps=None):
"""Generate a response conditioned on prompt via conditional masked diffusion.
1. Encode prompt as bytes
2. Create: [prompt_tokens | MASK * response_len]
3. Reverse diffusion: iteratively unmask response tokens
4. PC settling at each step refines predictions conditioned on prompt
"""
config = model.config
T = n_steps or config.n_diffusion_steps
mask_id = config.mask_token_id
prompt_tokens = encode_text(prompt_text)
prompt_len = len(prompt_tokens)
# Ensure we fit in max_seq_len
max_response = config.max_seq_len - prompt_len
if max_response <= 0:
prompt_tokens = prompt_tokens[:config.max_seq_len // 2]
prompt_len = len(prompt_tokens)
max_response = config.max_seq_len - prompt_len
response_len = min(response_len, max_response)
# Build initial sequence: [prompt | MASK...]
seq = prompt_tokens + [mask_id] * response_len
if len(seq) < config.max_seq_len:
seq = seq + [mask_id] * (config.max_seq_len - len(seq))
x = torch.tensor(seq[:config.max_seq_len], dtype=torch.long).unsqueeze(0).to(device)
# Response region mask (only unmask response, never touch prompt)
response_region = torch.zeros(1, config.max_seq_len, dtype=torch.bool, device=device)
response_region[0, prompt_len:prompt_len + response_len] = True
prev_h = None
prev_v = None
with torch.no_grad():
for t_val in range(T, 0, -1):
t = torch.full((1,), t_val, dtype=torch.long, device=device)
# Only consider masked tokens in response region
mask = (x == mask_id) & response_region
if not mask.any():
break
# Embed and settle
h_0 = model.embed_input(x, t)
h_init = model.amortized_forward_pass(h_0)
h_settled, v_final, _, _, _ = model.settle(
h_init, x, mask, t,
prev_h=prev_h, prev_v=prev_v,
)
# State continuation
prev_h = h_settled
prev_v = v_final
# Decode with temperature
logits = model.readout(model.readout_norm(h_settled[-1]))
if temperature > 0:
probs = F.softmax(logits / temperature, dim=-1)
else:
probs = F.softmax(logits, dim=-1)
# Determine how many tokens to unmask
t_prev = t_val - 1
target_mask_rate = model.schedule.mask_rate(
torch.tensor([t_prev], device=device)
).squeeze()
n_currently_masked = mask.sum().item()
n_to_keep_masked = int(target_mask_rate * response_len)
n_to_unmask = max(1, n_currently_masked - n_to_keep_masked)
# Score by confidence, only in response region
confidence = probs.max(dim=-1).values.squeeze() # (S,)
confidence[~mask.squeeze()] = float("inf")
# Unmask the most confident response tokens
masked_indices = mask.squeeze().nonzero(as_tuple=True)[0]
if len(masked_indices) > 0:
conf_at_masked = confidence[masked_indices]
n_to_unmask = min(n_to_unmask, len(masked_indices))
_, top_idx = conf_at_masked.topk(n_to_unmask)
unmask_pos = masked_indices[top_idx]
# Sample tokens
if temperature > 0:
for pos in unmask_pos:
token = torch.multinomial(probs[0, pos], 1).item()
x[0, pos] = token
else:
x[0, unmask_pos] = logits[0, unmask_pos].argmax(dim=-1)
# Final pass: unmask any remaining
remaining = (x == mask_id) & response_region
if remaining.any():
t_final = torch.ones(1, dtype=torch.long, device=device)
h_0 = model.embed_input(x, t_final)
h_init = model.amortized_forward_pass(h_0)
h_s, _, _, _, _ = model.settle(h_init, x, remaining, t_final, prev_h=prev_h, prev_v=prev_v)
logits = model.readout(model.readout_norm(h_s[-1]))
if temperature > 0:
for pos in remaining.squeeze().nonzero(as_tuple=True)[0]:
probs = F.softmax(logits[0, pos] / temperature, dim=-1)
x[0, pos] = torch.multinomial(probs, 1).item()
else:
x[0, remaining.squeeze()] = logits[0, remaining.squeeze()].argmax(dim=-1)
# Extract response
response_tokens = x[0, prompt_len:prompt_len + response_len].cpu().tolist()
return decode_tokens(response_tokens)
def main():
parser = argparse.ArgumentParser(description="PC-SHO-DLM Chat")
parser.add_argument("--checkpoint", type=str, default=None)
parser.add_argument("--device", type=str, default="auto")
parser.add_argument("--temperature", type=float, default=0.8)
parser.add_argument("--response_len", type=int, default=256)
args = parser.parse_args()
if args.device == "auto":
if torch.cuda.is_available(): args.device = "cuda"
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): args.device = "mps"
else: args.device = "cpu"
ckpt_path = args.checkpoint or find_checkpoint()
if not ckpt_path or not os.path.exists(ckpt_path):
print("No checkpoint found. Train first:")
print(" python3 train_chat.py --data alpaca")
sys.exit(1)
print(f"Loading {ckpt_path}...")
ckpt = torch.load(ckpt_path, map_location=args.device, weights_only=False)
config = ckpt["config"]
model = PCSHODLM(config).to(args.device)
model.load_state_dict(ckpt["model_state_dict"], strict=False)
model.eval()
print(f"Model: {sum(p.numel() for p in model.parameters()):,} params on {args.device}")
print()
print("=" * 60)
print(" PC-SHO-DLM Chat")
print(" Conditional Masked Diffusion Chatbot")
print("=" * 60)
print()
print(" Type a message and the model generates a response.")
print(" The prompt stays fixed; the response is iteratively")
print(" denoised via predictive-coding energy minimization.")
print()
print(" /temp 0.5 Set temperature (0=greedy, 1=creative)")
print(" /len 512 Set response length")
print(" /quit Exit")
print()
temp = args.temperature
resp_len = args.response_len
while True:
try:
user_input = input("You> ").strip()
except (EOFError, KeyboardInterrupt):
print("\nBye!")
break
if not user_input:
continue
if user_input.lower() in ["/quit", "/exit", "/q"]:
break
elif user_input.lower().startswith("/temp"):
try: temp = float(user_input.split()[1])
except: pass
print(f" Temperature: {temp}")
continue
elif user_input.lower().startswith("/len"):
try: resp_len = int(user_input.split()[1])
except: pass
print(f" Response length: {resp_len}")
continue
print(" [Generating...]")
response = generate_response(model, user_input, response_len=resp_len,
device=args.device, temperature=temp)
# Clean up: strip trailing nulls/garbage
response = response.rstrip('\x00').strip()
print(f"\nBot> {response}\n")
if __name__ == "__main__":
main()
|