File size: 3,866 Bytes
99b9cd2 | 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 | #!/usr/bin/env python3
"""Generate a planet name with the single-step ONNX model."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import numpy as np
import onnxruntime as ort
ROOT = Path(__file__).resolve().parent
DEFAULT_MODEL = ROOT / "planet_namer_fp16.onnx"
DEFAULT_VOCAB = ROOT / "vocab.json"
def generate_name(
stats: list[float],
*,
model_path: Path = DEFAULT_MODEL,
vocab_path: Path = DEFAULT_VOCAB,
temperature: float = 0.8,
seed: int | None = None,
) -> str:
"""Generate one name from seven normalized stats in vocab stat order."""
with vocab_path.open(encoding="utf-8") as handle:
vocab = json.load(handle)
stat_order = vocab["stat_order"]
if len(stats) != len(stat_order):
raise ValueError(
f"expected {len(stat_order)} stats ({', '.join(stat_order)}), "
f"received {len(stats)}"
)
if any(not 0.0 <= value <= 1.0 for value in stats):
raise ValueError("every stat must be between 0 and 1 inclusive")
if temperature < 0:
raise ValueError("temperature must be non-negative")
idx_to_char = {int(key): value for key, value in vocab["idx_to_char"].items()}
hidden_size = int(vocab.get("hidden_size", 192))
max_len = int(vocab.get("max_len", 20))
pad_idx = int(vocab.get("pad_idx", 0))
sos_idx = int(vocab.get("sos_idx", 1))
eos_idx = int(vocab.get("eos_idx", 2))
session = ort.InferenceSession(
str(model_path), providers=["CPUExecutionProvider"]
)
stats_array = np.asarray([stats], dtype=np.float32)
zero_stats = np.zeros_like(stats_array)
char_in = np.asarray([[sos_idx]], dtype=np.int64)
h_in = np.zeros((1, 1, hidden_size), dtype=np.float32)
c_in = np.zeros((1, 1, hidden_size), dtype=np.float32)
rng = np.random.default_rng(seed)
characters: list[str] = []
for step in range(max_len):
logits, h_in, c_in = session.run(
None,
{
"stats_init": stats_array if step == 0 else zero_stats,
"stats": stats_array,
"char_in": char_in,
"h_in": h_in,
"c_in": c_in,
},
)
if temperature < 1e-6:
char_idx = int(np.argmax(logits[0]))
else:
scaled = logits[0].astype(np.float64) / temperature
scaled -= scaled.max()
probabilities = np.exp(scaled)
probabilities /= probabilities.sum()
char_idx = int(rng.choice(len(probabilities), p=probabilities))
if char_idx in (pad_idx, eos_idx):
break
if char_idx != sos_idx:
character = idx_to_char.get(char_idx)
if character and not character.startswith("<"):
characters.append(character)
char_in = np.asarray([[char_idx]], dtype=np.int64)
return "".join(characters)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--stats",
nargs=7,
required=True,
type=float,
metavar=("ATM", "GRAV", "RES", "LIFE", "TEMP", "WATER", "RAD"),
help="seven normalized values in the documented order",
)
parser.add_argument("--temperature", type=float, default=0.8)
parser.add_argument("--seed", type=int)
parser.add_argument("--model", type=Path, default=DEFAULT_MODEL)
parser.add_argument("--vocab", type=Path, default=DEFAULT_VOCAB)
return parser.parse_args()
def main() -> None:
args = parse_args()
print(
generate_name(
args.stats,
model_path=args.model,
vocab_path=args.vocab,
temperature=args.temperature,
seed=args.seed,
)
)
if __name__ == "__main__":
main()
|