Update chat.py

chat.py

@@ -2,7 +2,8 @@
 """chat.py: ChatML chat with MetaDiffusion-600M checkpoints.
 
 Left-to-right block commit (semi-autoregressive): the leftmost masked
-positions are filled first, so <|im_end|> cannot win the race at position 0.
+positions are filled first, so <|im_end|> cannot win the race at position 0
+(which produced empty responses on the 150M architecture).
 
 Usage:
     Interactive:      python chat.py --model-path checkpoints/step_30000.pt
@@ -43,6 +44,15 @@ def load_model(model_path, device):
         sd = load_file(path / "model.safetensors")
         sd = {k[len("model."):] if k.startswith("model.") else k: v for k, v in sd.items()}
         model.load_state_dict(sd, strict=True)
+    elif path.suffix == ".safetensors":
+        from safetensors.torch import load_file
+        sd = load_file(model_path)
+        emb = sd["model.embed_tokens.weight"]
+        config = MetaDiffusionConfig(hidden_size=emb.shape[1],
+                                     mask_vocab_size=emb.shape[0])
+        model = MetaDiffusionLM(config).to(device)
+        sd = {k[len("model."):] if k.startswith("model.") else k: v for k, v in sd.items()}
+        model.load_state_dict(sd, strict=True)
     else:
         ckpt = torch.load(model_path, map_location=device, weights_only=False)
         config = build_config(ckpt["config"])
@@ -57,7 +67,7 @@ def load_model(model_path, device):
 
 
 def ensure_special_tokens(tokenizer):
-    """Add [MASK] + rainbow if missing"""
+    """Add [MASK] + rainbow if missing (source tokenizer case)."""
     added = []
     if tokenizer.convert_tokens_to_ids("[MASK]") == tokenizer.unk_token_id:
         added.append("[MASK]")