| """Word-level tokenizer. Unknown words map to id 0 (<pad>), which the |
| embedding layer ignores, so garbage prompts are harmless.""" |
|
|
| import re |
|
|
| import torch |
|
|
| from .palette import COLOR_WORDS |
| from .templates import CLASSES |
|
|
| CLASS_SYNONYMS = { |
| "potion": ["potion", "bottle", "flask", "elixir", "vial"], |
| "sword": ["sword", "blade", "dagger"], |
| "shield": ["shield", "buckler"], |
| "ship": ["ship", "rocket", "spaceship", "starship"], |
| "invader": ["invader", "alien", "monster"], |
| "slime": ["slime", "blob", "jelly"], |
| "skull": ["skull", "bone", "skeleton"], |
| "tree": ["tree", "pine", "forest"], |
| "key": ["key"], |
| "heart": ["heart", "love", "life"], |
| "star": ["star", "sparkle"], |
| "mushroom": ["mushroom", "fungus", "toadstool", "shroom"], |
| "ghost": ["ghost", "spirit", "phantom", "spooky"], |
| "bomb": ["bomb", "grenade", "explosive"], |
| "coin": ["coin", "token", "money", "medal"], |
| "gem": ["gem", "diamond", "jewel", "crystal"], |
| "cherry": ["cherry", "cherries", "berry", "fruit"], |
| "crown": ["crown"], |
| "bat": ["bat", "vampire"], |
| "cactus": ["cactus", "desert", "succulent"], |
| "robot": ["robot", "droid", "android", "mech"], |
| } |
|
|
| PAD = "<pad>" |
|
|
| |
| ACCESSORY_SYNONYMS = { |
| "none": [], |
| "crown": ["crown", "king", "queen", "royal", "crowned", "prince", "princess"], |
| "halo": ["halo", "angel", "angelic", "holy", "saint"], |
| "horns": ["horns", "horned", "devil", "demon", "evil"], |
| "helmet": ["helmet", "knight", "armor", "armour", "armored"], |
| "antenna": ["antenna", "antennae", "bug", "martian"], |
| "wings": ["wings", "winged", "fly", "flying", "bird", "fairy", "butterfly"], |
| "torch": ["torch", "torches", "fire", "flame", "burning"], |
| } |
| SIZE_WORDS = {"small": 1, "tiny": 1, "mini": 1, "little": 1} |
| |
| SEPARATORS = ["with", "wearing", "holding", "and", "plus"] |
|
|
|
|
| def build_vocab(): |
| vocab = {PAD: 0} |
| groups = [ |
| set(w for ws in CLASS_SYNONYMS.values() for w in ws), |
| set(COLOR_WORDS), |
| set(w for ws in ACCESSORY_SYNONYMS.values() for w in ws), |
| set(SIZE_WORDS), |
| set(SEPARATORS), |
| ] |
| for g in groups: |
| for word in sorted(g): |
| if word not in vocab: |
| vocab[word] = len(vocab) |
| return vocab |
|
|
|
|
| VOCAB = build_vocab() |
|
|
| _SPLIT = re.compile(r"[^a-z]+") |
|
|
|
|
| def tokenize(prompt, vocab=VOCAB): |
| words = [w for w in _SPLIT.split(prompt.lower()) if w] |
| ids = [vocab.get(w, 0) for w in words] |
| return torch.tensor(ids if ids else [0], dtype=torch.long) |
|
|
|
|
| def pad_batch(prompts, vocab=VOCAB): |
| toks = [tokenize(p, vocab) for p in prompts] |
| L = max(t.numel() for t in toks) |
| out = torch.zeros(len(toks), L, dtype=torch.long) |
| for i, t in enumerate(toks): |
| out[i, : t.numel()] = t |
| return out |
|
|
|
|
| def split_prompt(prompt): |
| """'green slime with gold crown' -> ('green slime', 'gold crown'). |
| No separator -> whole prompt is the body half, accessory half empty.""" |
| words = [w for w in _SPLIT.split(prompt.lower()) if w] |
| for i, w in enumerate(words): |
| if w in SEPARATORS: |
| return " ".join(words[:i]), " ".join(words[i + 1:]) |
| return " ".join(words), "" |
|
|
|
|
| def pad_batch_split(prompts, vocab=VOCAB): |
| halves = [split_prompt(p) for p in prompts] |
| return pad_batch([a for a, _ in halves], vocab), pad_batch([b for _, b in halves], vocab) |
|
|
|
|
| def word_class(word): |
| for cls, ws in CLASS_SYNONYMS.items(): |
| if word in ws: |
| return CLASSES.index(cls) |
| return None |
|
|
