feat: Initial Commit

model/builder_model.py

@@ -0,0 +1,133 @@
+import torch
+import torch.nn as nn
+
+from config import (
+    DEVICE,
+    START_TOKEN,
+    END_TOKEN,
+    MAX_OUTPUT_LENGTH,
+)
+
+from model.transformer import MinecraftTransformer
+
+
+class BuilderModel(nn.Module):
+    """
+    High-level wrapper around the MinecraftTransformer.
+    """
+
+    def __init__(self, tokenizer):
+        super().__init__()
+
+        self.tokenizer = tokenizer
+
+        self.transformer = MinecraftTransformer(
+            tokenizer.text_vocab_size,
+            tokenizer.output_vocab_size,
+        )
+
+    def forward(self, src, tgt):
+        """
+        Training forward pass.
+
+        src = prompt tokens
+        tgt = decoder input tokens (already shifted by train.py)
+        """
+
+        return self.transformer(
+            src=src,
+            tgt=tgt,
+        )
+
+    @torch.no_grad()
+    def generate(
+            self,
+            prompt,
+            max_length=MAX_OUTPUT_LENGTH,
+    ):
+        """
+        Generate a Minecraft structure from a prompt.
+        """
+
+        self.eval()
+
+        device = next(self.parameters()).device
+
+        # Encode prompt
+        src = torch.tensor(
+            [self.tokenizer.encode_prompt(prompt)],
+            dtype=torch.long,
+            device=device,
+        )
+
+        # Encode with transformer
+        memory = self.transformer.encode(src)
+
+        start_id = self.tokenizer.output_to_id[START_TOKEN]
+        end_id = self.tokenizer.output_to_id[END_TOKEN]
+
+        generated = [start_id]
+
+        for _ in range(max_length):
+
+            tgt = torch.tensor(
+                [generated],
+                dtype=torch.long,
+                device=device,
+            )
+
+            logits = self.transformer.decode(
+                tgt,
+                memory,
+            )
+
+            probs = torch.softmax(logits[0, -1], dim=-1)
+
+            next_token = torch.argmax(probs).item()
+
+            print(
+                next_token,
+                self.tokenizer.id_to_output[next_token],
+                f"{probs[next_token].item():.3f}"
+            )
+
+            generated.append(next_token)
+
+            if next_token == end_id:
+                break
+
+        blocks = self.tokenizer.decode_blocks(generated)
+
+        # Don't count the <START> token
+        token_count = max(0, len(generated) - 1)
+
+        return blocks, token_count
+
+    def save(self, path):
+        """
+        Save model weights.
+        """
+
+        torch.save(
+            self.state_dict(),
+            path,
+        )
+
+    def load(self, path):
+        """
+        Load model weights.
+        """
+
+        self.load_state_dict(
+            torch.load(
+                path,
+                map_location=DEVICE,
+            )
+        )
+
+    def to_device(self):
+        """
+        Move model to configured device.
+        """
+
+        return self.to(DEVICE)