133 lines
2.8 KiB
Python
133 lines
2.8 KiB
Python
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) |