feat: Initial Commit
This commit is contained in:
@@ -0,0 +1,655 @@
|
||||
[
|
||||
{
|
||||
"prompt": "single stone",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:stone"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single oak planks",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_planks"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single oak log",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_log"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single spruce planks",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:spruce_planks"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single birch planks",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:birch_planks"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single jungle planks",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:jungle_planks"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single acacia planks",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:acacia_planks"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single dark oak planks",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:dark_oak_planks"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single mangrove planks",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:mangrove_planks"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single cherry planks",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:cherry_planks"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single bamboo planks",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:bamboo_planks"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single crimson planks",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:crimson_planks"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single warped planks",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:warped_planks"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single oak log",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_log"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single spruce log",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:spruce_log"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single birch log",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:birch_log"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single jungle log",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:jungle_log"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single acacia log",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:acacia_log"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single dark oak log",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:dark_oak_log"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single mangrove log",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:mangrove_log"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single cherry log",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:cherry_log"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single crimson stem",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:crimson_stem"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "single warped stem",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:warped_stem"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "3 block tall wooden pillar",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_log"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_log"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_log"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "3 block tall wooden pillar out of oak logs",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_log"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_log"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_log"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "3 block tall wooden pillar out of spruce logs",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:spruce_log"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 0,
|
||||
"block": "minecraft:spruce_log"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 0,
|
||||
"block": "minecraft:spruce_log"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "3 block tall wooden pillar out of birch logs",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:birch_log"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 0,
|
||||
"block": "minecraft:birch_log"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 0,
|
||||
"block": "minecraft:birch_log"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "a 4 block wide and 3 block tall wooden wall",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 1,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 2,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 3,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 1,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 2,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 3,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 1,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 2,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 3,
|
||||
"block": "minecraft:oak_planks"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "a 4 block wide and 3 block tall wooden wall out of oak planks",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 1,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 2,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 3,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 1,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 2,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 3,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 1,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 2,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 3,
|
||||
"block": "minecraft:oak_planks"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "a small wooden wall",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 1,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 2,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 3,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 1,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 2,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 3,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 1,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 2,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 3,
|
||||
"block": "minecraft:oak_planks"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"prompt": "a small wooden wall out of oak planks",
|
||||
"blocks": [
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 1,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 2,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z": 3,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 1,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 2,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"z": 3,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 0,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 1,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 2,
|
||||
"block": "minecraft:oak_planks"
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"z": 3,
|
||||
"block": "minecraft:oak_planks"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
+113
@@ -0,0 +1,113 @@
|
||||
import json
|
||||
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from torch.nn.utils.rnn import pad_sequence
|
||||
|
||||
from config import (
|
||||
PAD_TOKEN,
|
||||
MAX_PROMPT_LENGTH,
|
||||
MAX_OUTPUT_LENGTH,
|
||||
)
|
||||
|
||||
|
||||
class MinecraftDataset(Dataset):
|
||||
"""
|
||||
Dataset for training the Minecraft Builder AI.
|
||||
"""
|
||||
|
||||
def __init__(self, dataset_path, tokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
with open(dataset_path, "r", encoding="utf8") as f:
|
||||
self.data = json.load(f)
|
||||
|
||||
self.pad_text = tokenizer.text_to_id[PAD_TOKEN]
|
||||
self.pad_output = tokenizer.output_to_id[PAD_TOKEN]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data)
|
||||
|
||||
def __getitem__(self, index):
|
||||
sample = self.data[index]
|
||||
|
||||
prompt_ids = self.tokenizer.encode_prompt(
|
||||
sample["prompt"]
|
||||
)
|
||||
|
||||
output_ids = self.tokenizer.encode_blocks(
|
||||
sample["blocks"]
|
||||
)
|
||||
|
||||
# Limit maximum sequence lengths
|
||||
prompt_ids = prompt_ids[:MAX_PROMPT_LENGTH]
|
||||
output_ids = output_ids[:MAX_OUTPUT_LENGTH]
|
||||
|
||||
return (
|
||||
torch.tensor(prompt_ids, dtype=torch.long),
|
||||
torch.tensor(output_ids, dtype=torch.long),
|
||||
)
|
||||
|
||||
|
||||
def collate_fn(batch):
|
||||
"""
|
||||
Pads sequences inside a batch.
|
||||
"""
|
||||
|
||||
prompts = [item[0] for item in batch]
|
||||
outputs = [item[1] for item in batch]
|
||||
|
||||
prompt_pad = prompts[0].new_tensor(
|
||||
[0]
|
||||
) # placeholder (overwritten below)
|
||||
|
||||
output_pad = outputs[0].new_tensor(
|
||||
[0]
|
||||
)
|
||||
|
||||
# These values are replaced by the DataLoader factory below.
|
||||
prompt_padding_value = getattr(collate_fn, "prompt_pad", 0)
|
||||
output_padding_value = getattr(collate_fn, "output_pad", 0)
|
||||
|
||||
prompts = pad_sequence(
|
||||
prompts,
|
||||
batch_first=True,
|
||||
padding_value=prompt_padding_value,
|
||||
)
|
||||
|
||||
outputs = pad_sequence(
|
||||
outputs,
|
||||
batch_first=True,
|
||||
padding_value=output_padding_value,
|
||||
)
|
||||
|
||||
return prompts, outputs
|
||||
|
||||
|
||||
def create_dataloader(
|
||||
dataset_path,
|
||||
tokenizer,
|
||||
batch_size,
|
||||
shuffle=True,
|
||||
):
|
||||
"""
|
||||
Creates a DataLoader with automatic padding.
|
||||
"""
|
||||
|
||||
dataset = MinecraftDataset(
|
||||
dataset_path,
|
||||
tokenizer,
|
||||
)
|
||||
|
||||
# Give the collate function the correct padding IDs
|
||||
collate_fn.prompt_pad = tokenizer.text_to_id[PAD_TOKEN]
|
||||
collate_fn.output_pad = tokenizer.output_to_id[PAD_TOKEN]
|
||||
|
||||
loader = torch.utils.data.DataLoader(
|
||||
dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=shuffle,
|
||||
collate_fn=collate_fn,
|
||||
)
|
||||
|
||||
return loader
|
||||
@@ -0,0 +1,232 @@
|
||||
import json
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
|
||||
from config import (
|
||||
PAD_TOKEN,
|
||||
START_TOKEN,
|
||||
END_TOKEN,
|
||||
UNK_TOKEN,
|
||||
SPECIAL_TOKENS,
|
||||
)
|
||||
|
||||
|
||||
class Tokenizer:
|
||||
def __init__(self):
|
||||
self.text_to_id = {}
|
||||
self.id_to_text = {}
|
||||
|
||||
self.output_to_id = {}
|
||||
self.id_to_output = {}
|
||||
|
||||
# =====================================================
|
||||
# Build Vocabulary
|
||||
# =====================================================
|
||||
|
||||
def build(self, dataset_path):
|
||||
|
||||
dataset_path = Path(dataset_path)
|
||||
|
||||
with open(dataset_path, "r", encoding="utf8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
text_counter = Counter()
|
||||
output_counter = Counter()
|
||||
|
||||
for sample in data:
|
||||
|
||||
# ---------------------------
|
||||
# Prompt tokens
|
||||
# ---------------------------
|
||||
|
||||
prompt = sample["prompt"].lower().split()
|
||||
|
||||
text_counter.update(prompt)
|
||||
|
||||
# ---------------------------
|
||||
# Output tokens
|
||||
# ---------------------------
|
||||
|
||||
for block in sample["blocks"]:
|
||||
|
||||
x = block["x"]
|
||||
y = block["y"]
|
||||
z = block["z"]
|
||||
b = block["block"]
|
||||
|
||||
output_counter.update([
|
||||
f"X_{x}",
|
||||
f"Y_{y}",
|
||||
f"Z_{z}",
|
||||
f"BLOCK_{b}"
|
||||
])
|
||||
|
||||
# Special tokens
|
||||
|
||||
text_vocab = SPECIAL_TOKENS + sorted(text_counter.keys())
|
||||
output_vocab = SPECIAL_TOKENS + sorted(output_counter.keys())
|
||||
|
||||
self.text_to_id = {
|
||||
token: i for i, token in enumerate(text_vocab)
|
||||
}
|
||||
|
||||
self.id_to_text = {
|
||||
i: token for token, i in self.text_to_id.items()
|
||||
}
|
||||
|
||||
self.output_to_id = {
|
||||
token: i for i, token in enumerate(output_vocab)
|
||||
}
|
||||
|
||||
self.id_to_output = {
|
||||
i: token for token, i in self.output_to_id.items()
|
||||
}
|
||||
|
||||
# =====================================================
|
||||
# Prompt Encoding
|
||||
# =====================================================
|
||||
|
||||
def encode_prompt(self, prompt):
|
||||
|
||||
tokens = prompt.lower().split()
|
||||
|
||||
ids = [
|
||||
self.text_to_id[START_TOKEN]
|
||||
]
|
||||
|
||||
for token in tokens:
|
||||
|
||||
ids.append(
|
||||
self.text_to_id.get(
|
||||
token,
|
||||
self.text_to_id[UNK_TOKEN]
|
||||
)
|
||||
)
|
||||
|
||||
ids.append(
|
||||
self.text_to_id[END_TOKEN]
|
||||
)
|
||||
|
||||
return ids
|
||||
|
||||
def decode_prompt(self, ids):
|
||||
|
||||
words = []
|
||||
|
||||
for idx in ids:
|
||||
|
||||
token = self.id_to_text[idx]
|
||||
|
||||
if token in SPECIAL_TOKENS:
|
||||
continue
|
||||
|
||||
words.append(token)
|
||||
|
||||
return " ".join(words)
|
||||
|
||||
# =====================================================
|
||||
# Structure Encoding
|
||||
# =====================================================
|
||||
|
||||
def encode_blocks(self, blocks):
|
||||
|
||||
ids = [
|
||||
self.output_to_id[START_TOKEN]
|
||||
]
|
||||
|
||||
for block in blocks:
|
||||
|
||||
ids.extend([
|
||||
self.output_to_id[f"X_{block['x']}"],
|
||||
self.output_to_id[f"Y_{block['y']}"],
|
||||
self.output_to_id[f"Z_{block['z']}"],
|
||||
self.output_to_id[f"BLOCK_{block['block']}"],
|
||||
])
|
||||
|
||||
ids.append(
|
||||
self.output_to_id[END_TOKEN]
|
||||
)
|
||||
|
||||
return ids
|
||||
|
||||
def decode_blocks(self, ids):
|
||||
|
||||
tokens = []
|
||||
|
||||
for idx in ids:
|
||||
|
||||
token = self.id_to_output[idx]
|
||||
|
||||
if token in SPECIAL_TOKENS:
|
||||
continue
|
||||
|
||||
tokens.append(token)
|
||||
|
||||
blocks = []
|
||||
|
||||
i = 0
|
||||
|
||||
while i + 3 < len(tokens):
|
||||
|
||||
x = int(tokens[i][2:])
|
||||
y = int(tokens[i + 1][2:])
|
||||
z = int(tokens[i + 2][2:])
|
||||
block = tokens[i + 3][6:]
|
||||
|
||||
blocks.append({
|
||||
"x": x,
|
||||
"y": y,
|
||||
"z": z,
|
||||
"block": block
|
||||
})
|
||||
|
||||
i += 4
|
||||
|
||||
return blocks
|
||||
|
||||
# =====================================================
|
||||
# Save / Load
|
||||
# =====================================================
|
||||
|
||||
def save(self, folder):
|
||||
|
||||
folder = Path(folder)
|
||||
folder.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with open(folder / "text_vocab.json", "w") as f:
|
||||
json.dump(self.text_to_id, f, indent=4)
|
||||
|
||||
with open(folder / "output_vocab.json", "w") as f:
|
||||
json.dump(self.output_to_id, f, indent=4)
|
||||
|
||||
def load(self, folder):
|
||||
|
||||
folder = Path(folder)
|
||||
|
||||
with open(folder / "text_vocab.json", "r") as f:
|
||||
self.text_to_id = json.load(f)
|
||||
|
||||
with open(folder / "output_vocab.json", "r") as f:
|
||||
self.output_to_id = json.load(f)
|
||||
|
||||
self.id_to_text = {
|
||||
int(v): k
|
||||
for k, v in self.text_to_id.items()
|
||||
}
|
||||
|
||||
self.id_to_output = {
|
||||
int(v): k
|
||||
for k, v in self.output_to_id.items()
|
||||
}
|
||||
|
||||
# =====================================================
|
||||
# Properties
|
||||
# =====================================================
|
||||
|
||||
@property
|
||||
def text_vocab_size(self):
|
||||
return len(self.text_to_id)
|
||||
|
||||
@property
|
||||
def output_vocab_size(self):
|
||||
return len(self.output_to_id)
|
||||
Reference in New Issue
Block a user