File size: 4,515 Bytes
ed2140a | 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 | """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()
|