Upload inference.py with huggingface_hub

inference.py

@@ -0,0 +1,136 @@
+"""Inference script for Diffusion-LM Riddle Solver (Hugging Face model).
+
+Usage:
+    python3 inference.py --riddle "i speak without a mouth what am i"
+"""
+
+import json
+import sys
+import argparse
+import warnings
+
+
+def load_model(model_dir: str = "."):
+    """Load model weights and config from HF model directory."""
+    import torch
+    
+    with open(f"{model_dir}/config.json") as f:
+        config = json.load(f)
+    
+    with open(f"{model_dir}/vocab.json") as f:
+        vocab = json.load(f)
+    
+    inv_vocab = {int(v): k for k, v in vocab.items()}
+    
+    # Add riddle_diffusion.py to Python path if running from HF directory
+    sys.path.insert(0, model_dir)
+    from riddle_diffusion import DiffusionRiddleModel
+    from riddle_diffusion import get_schedule
+    
+    model = DiffusionRiddleModel(
+        vocab_size=config["vocab_size"],
+        d_model=config["d_model"],
+        n_layers=config["n_layers"],
+        d_ff=config["d_ff"],
+        n_heads=config["n_heads"],
+        a_len=config["a_len"],
+        q_len=config["q_len"],
+        T=config["T"],
+    )
+    
+    state = torch.load(f"{model_dir}/model.safetensors", map_location="cpu",
+                       weights_only=True)
+    model.load_state_dict(state)
+    model.eval()
+    
+    return model, config, vocab, inv_vocab
+
+
+def predict(model, config, vocab, inv_vocab, riddle: str, k_samples: int = 10,
+            device: str = "cpu"):
+    """Run prediction on a single riddle."""
+    import torch
+    import torch.nn.functional as F
+    
+    model.to(device)
+    
+    # Tokenize
+    tokens = [vocab.get(w, vocab.get("<UNK>", 1)) for w in riddle.lower().split()]
+    if len(tokens) > config["q_len"]:
+        tokens = tokens[:config["q_len"]]
+    
+    q_tokens = torch.tensor([tokens], device=device)
+    
+    # Diffusion schedule (sqrt power law)
+    betas = torch.sqrt(torch.linspace(1e-4, 0.02, config["T"])).to(device)
+    alphas = 1.0 - betas
+    alpha_bars = torch.cumprod(alphas, dim=0)
+    
+    # Reverse diffusion
+    x_t = torch.randn(k_samples, config["a_len"], config["d_model"], device=device)
+    
+    for t in reversed(range(config["T"])):
+        t_tensor = torch.full((k_samples,), t, device=device, dtype=torch.long)
+        pred_x0 = model(x_t, t_tensor, q_tokens)
+        
+        # Euclidean clamping
+        logits = 2.0 * F.linear(pred_x0, model.emb.weight)
+        logits = logits - model.emb.weight.square().sum(dim=-1).unsqueeze(0).unsqueeze(0)
+        x0_tokens = logits.argmax(dim=-1)
+        x0_emb = model.emb(x0_tokens)
+        
+        if t > 0:
+            alpha_bar = alpha_bars[t]
+            alpha_bar_prev = alpha_bars[t - 1]
+            beta_tilde = betas[t] * (1 - alpha_bar_prev) / (1 - alpha_bar)
+            noise = torch.randn_like(x_t)
+            coef1 = torch.sqrt(alpha_bar_prev) * betas[t] / (1 - alpha_bar)
+            coef2 = torch.sqrt(alpha_bar) * (1 - alpha_bar_prev) / (1 - alpha_bar)
+            mu = coef1 * x0_emb + coef2 * x_t
+            x_t = mu + torch.sqrt(beta_tilde) * noise
+        else:
+            x_t = x0_emb
+    
+    # Decode
+    pred_tokens = []
+    for b in range(k_samples):
+        pred = ""
+        for pos in range(config["a_len"]):
+            tok_id = x0_tokens[b, pos].item()
+            if tok_id == 0:
+                break
+            pred += inv_vocab.get(tok_id, "?") + " "
+        pred_tokens.append(pred.strip())
+    
+    # Majority vote
+    from collections import Counter
+    counts = Counter(pred_tokens)
+    winner = counts.most_common(1)[0][0]
+    
+    return winner, pred_tokens
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--riddle", required=True, help="Riddle text")
+    parser.add_argument("--model-dir", default=".", help="Model directory")
+    parser.add_argument("--device", default="cpu", help="Device (cpu, mps, cuda)")
+    parser.add_argument("--k-samples", type=int, default=10)
+    args = parser.parse_args()
+    
+    model, config, vocab, inv_vocab = load_model(args.model_dir)
+    answer, candidates = predict(model, config, vocab, inv_vocab,
+                                 args.riddle, args.k_samples, args.device)
+    
+    print(f"Riddle:  {args.riddle}")
+    print(f"Answer:  {answer}")
+    if len(set(candidates)) > 1:
+        from collections import Counter
+        counts = Counter(candidates)
+        print(f"Candidates ({args.k_samples} samples):")
+        for text, count in counts.most_common():
+            print(f"  {text:<20} ({count} votes)")
+
+
+if __name__ == "__main__":
+    main()