feat: Initial Commit
This commit is contained in:
Generated
+10
@@ -0,0 +1,10 @@
|
||||
# Default ignored files
|
||||
/shelf/
|
||||
/workspace.xml
|
||||
# Editor-based HTTP Client requests
|
||||
/httpRequests/
|
||||
# Ignored default folder with query files
|
||||
/queries/
|
||||
# Datasource local storage ignored files
|
||||
/dataSources/
|
||||
/dataSources.local.xml
|
||||
Generated
+9
@@ -0,0 +1,9 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="Black">
|
||||
<option name="sdkName" value="uv (MineBuildAI)" />
|
||||
</component>
|
||||
<component name="ProjectRootManager" version="2" project-jdk-name="uv (MineBuildAI)" project-jdk-type="Python SDK">
|
||||
<output url="file://$PROJECT_DIR$/out" />
|
||||
</component>
|
||||
</project>
|
||||
Generated
+8
@@ -0,0 +1,8 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="ProjectModuleManager">
|
||||
<modules>
|
||||
<module fileurl="file://$PROJECT_DIR$/MineBuildAI.iml" filepath="$PROJECT_DIR$/MineBuildAI.iml" />
|
||||
</modules>
|
||||
</component>
|
||||
</project>
|
||||
Generated
+13
@@ -0,0 +1,13 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="CommitMessageInspectionProfile">
|
||||
<profile version="1.0">
|
||||
<inspection_tool class="CommitFormat" enabled="true" level="WARNING" enabled_by_default="true" />
|
||||
<inspection_tool class="CommitMessageSpellChecking" enabled="true" level="TYPO" enabled_by_default="true" />
|
||||
<inspection_tool class="CommitNamingConvention" enabled="true" level="WARNING" enabled_by_default="true" />
|
||||
</profile>
|
||||
</component>
|
||||
<component name="VcsDirectoryMappings">
|
||||
<mapping directory="$PROJECT_DIR$" vcs="Git" />
|
||||
</component>
|
||||
</project>
|
||||
@@ -0,0 +1,11 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<module type="PYTHON_MODULE" version="4">
|
||||
<component name="NewModuleRootManager" inherit-compiler-output="true">
|
||||
<exclude-output />
|
||||
<content url="file://$MODULE_DIR$">
|
||||
<excludeFolder url="file://$MODULE_DIR$/.venv" />
|
||||
</content>
|
||||
<orderEntry type="jdk" jdkName="uv (MineBuildAI)" jdkType="Python SDK" />
|
||||
<orderEntry type="sourceFolder" forTests="false" />
|
||||
</component>
|
||||
</module>
|
||||
@@ -0,0 +1,123 @@
|
||||
"""
|
||||
config.py
|
||||
|
||||
Global configuration for the Minecraft AI Builder.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
import torch
|
||||
|
||||
# ==========================================================
|
||||
# Project Paths
|
||||
# ==========================================================
|
||||
|
||||
ROOT_DIR = Path(__file__).parent
|
||||
|
||||
DATASET_PATH = ROOT_DIR / "data" / "dataset.json"
|
||||
CHECKPOINT_DIR = ROOT_DIR / "checkpoints"
|
||||
|
||||
CHECKPOINT_DIR.mkdir(exist_ok=True)
|
||||
|
||||
# ==========================================================
|
||||
# Model
|
||||
# ==========================================================
|
||||
|
||||
EMBED_DIM = 512
|
||||
NUM_HEADS = 8
|
||||
NUM_ENCODER_LAYERS = 6
|
||||
NUM_DECODER_LAYERS = 6
|
||||
FEED_FORWARD_DIM = 2048
|
||||
DROPOUT = 0.1
|
||||
|
||||
# Maximum lengths
|
||||
|
||||
MAX_PROMPT_LENGTH = 64
|
||||
MAX_OUTPUT_LENGTH = 4096
|
||||
|
||||
# ==========================================================
|
||||
# Vocabulary
|
||||
# ==========================================================
|
||||
|
||||
# Filled automatically after building vocab
|
||||
TEXT_VOCAB_SIZE = None
|
||||
OUTPUT_VOCAB_SIZE = None
|
||||
|
||||
# Special Tokens
|
||||
PAD_TOKEN = "<PAD>"
|
||||
START_TOKEN = "<START>"
|
||||
END_TOKEN = "<END>"
|
||||
UNK_TOKEN = "<UNK>"
|
||||
|
||||
SPECIAL_TOKENS = [
|
||||
PAD_TOKEN,
|
||||
START_TOKEN,
|
||||
END_TOKEN,
|
||||
UNK_TOKEN
|
||||
]
|
||||
|
||||
# ==========================================================
|
||||
# Training
|
||||
# ==========================================================
|
||||
|
||||
BATCH_SIZE = 8
|
||||
LEARNING_RATE = 1e-4
|
||||
|
||||
EPOCHS = 100
|
||||
|
||||
WEIGHT_DECAY = 1e-5
|
||||
|
||||
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
# ==========================================================
|
||||
# Generation
|
||||
# ==========================================================
|
||||
|
||||
TEMPERATURE = 1.0
|
||||
|
||||
TOP_K = 40
|
||||
|
||||
TOP_P = 0.95
|
||||
|
||||
# ==========================================================
|
||||
# Random Seed
|
||||
# ==========================================================
|
||||
|
||||
RANDOM_SEED = 42
|
||||
|
||||
# ==========================================================
|
||||
# Coordinate Limits
|
||||
# ==========================================================
|
||||
|
||||
MIN_COORD = 0
|
||||
MAX_COORD = 128
|
||||
|
||||
# ==========================================================
|
||||
# Supported Minecraft Blocks
|
||||
# ==========================================================
|
||||
|
||||
SUPPORTED_BLOCKS = [
|
||||
"minecraft:stone",
|
||||
|
||||
"minecraft:oak_planks",
|
||||
"minecraft:spruce_planks",
|
||||
"minecraft:birch_planks",
|
||||
"minecraft:jungle_planks",
|
||||
"minecraft:acacia_planks",
|
||||
"minecraft:dark_oak_planks",
|
||||
"minecraft:mangrove_planks",
|
||||
"minecraft:cherry_planks",
|
||||
"minecraft:bamboo_planks",
|
||||
"minecraft:crimson_planks",
|
||||
"minecraft:warped_planks",
|
||||
|
||||
"minecraft:oak_log",
|
||||
"minecraft:spruce_log",
|
||||
"minecraft:birch_log",
|
||||
"minecraft:jungle_log",
|
||||
"minecraft:acacia_log",
|
||||
"minecraft:dark_oak_log",
|
||||
"minecraft:mangrove_log",
|
||||
"minecraft:cherry_log",
|
||||
"minecraft:crimson_stem",
|
||||
"minecraft:warped_stem",
|
||||
]
|
||||
@@ -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)
|
||||
@@ -0,0 +1,74 @@
|
||||
from config import (
|
||||
DEVICE,
|
||||
CHECKPOINT_DIR,
|
||||
)
|
||||
|
||||
from data.tokenizer import Tokenizer
|
||||
from model.builder_model import BuilderModel
|
||||
from utils.checkpoint import latest_checkpoint, load_checkpoint
|
||||
|
||||
|
||||
def main():
|
||||
if DEVICE == "cuda":
|
||||
print("Cuda Mode: true")
|
||||
else:
|
||||
print("Cuda Mode: false")
|
||||
|
||||
print("Loading tokenizer...")
|
||||
|
||||
tokenizer = Tokenizer()
|
||||
tokenizer.load("data")
|
||||
|
||||
print("Loading model...")
|
||||
|
||||
model = BuilderModel(tokenizer).to(DEVICE)
|
||||
|
||||
checkpoint = latest_checkpoint(CHECKPOINT_DIR)
|
||||
|
||||
if checkpoint is None:
|
||||
print("No checkpoint found!")
|
||||
return
|
||||
|
||||
print(f"Loading checkpoint: {checkpoint}")
|
||||
|
||||
load_checkpoint(
|
||||
checkpoint,
|
||||
model,
|
||||
optimizer=None,
|
||||
map_location=DEVICE,
|
||||
)
|
||||
|
||||
model.eval()
|
||||
|
||||
print("\nMinecraft AI Builder")
|
||||
print("Type 'exit' to quit.\n")
|
||||
|
||||
while True:
|
||||
|
||||
prompt = input("> ")
|
||||
|
||||
if prompt.lower() in ("exit", "quit"):
|
||||
break
|
||||
|
||||
blocks, token_count = model.generate(prompt)
|
||||
|
||||
print("\nGenerated Structure:\n")
|
||||
|
||||
if len(blocks) == 0:
|
||||
print("(No blocks generated)")
|
||||
continue
|
||||
|
||||
for block in blocks:
|
||||
print(
|
||||
f"{block['block']:20}"
|
||||
f" x={block['x']:4}"
|
||||
f" y={block['y']:4}"
|
||||
f" z={block['z']:4}"
|
||||
)
|
||||
|
||||
print(f"\nTotal blocks: {len(blocks)}")
|
||||
print(f"Generated tokens: {token_count}\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,53 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def json_string_to_py(json_text: str, output_file: str = "generated.py"):
|
||||
"""
|
||||
Convert an AI response containing:
|
||||
{
|
||||
"file_name": "...",
|
||||
"script": "..."
|
||||
}
|
||||
|
||||
into a valid Python .py file.
|
||||
"""
|
||||
|
||||
# Parse the JSON response
|
||||
data = json.loads(json_text)
|
||||
|
||||
# Validate that the script exists
|
||||
if "script" not in data:
|
||||
raise ValueError("JSON does not contain a 'script' field.")
|
||||
|
||||
script = data["script"]
|
||||
|
||||
if not isinstance(script, str):
|
||||
raise TypeError("'script' must be a string.")
|
||||
|
||||
# Write the Python source code exactly as provided
|
||||
output_path = Path("builders/" + output_file)
|
||||
output_path.write_text(script, encoding="utf-8")
|
||||
|
||||
print(f"Created: {output_path.resolve()}")
|
||||
|
||||
return output_path
|
||||
|
||||
|
||||
# Example usage
|
||||
if __name__ == "__main__":
|
||||
|
||||
ai_response = r'''
|
||||
{
|
||||
"file_name": "small_oak_tree",
|
||||
"script": "import random\n\nSEED = 42\nrandom.seed(SEED)\n\nprint(\"Hello, World!\")"
|
||||
}
|
||||
'''
|
||||
|
||||
data = json.loads(ai_response)
|
||||
|
||||
# Prefer the AI-provided filename
|
||||
filename = data.get("file_name", "generated")
|
||||
output_file = f"{filename}.py"
|
||||
|
||||
json_string_to_py(ai_response, output_file)
|
||||
@@ -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)
|
||||
@@ -0,0 +1,55 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class PositionalEncoding(nn.Module):
|
||||
"""
|
||||
Standard sinusoidal positional encoding from
|
||||
"Attention Is All You Need".
|
||||
"""
|
||||
|
||||
def __init__(self, embed_dim, dropout=0.1, max_len=10000):
|
||||
super().__init__()
|
||||
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
pe = torch.zeros(max_len, embed_dim)
|
||||
|
||||
position = torch.arange(
|
||||
0,
|
||||
max_len,
|
||||
dtype=torch.float
|
||||
).unsqueeze(1)
|
||||
|
||||
div_term = torch.exp(
|
||||
torch.arange(
|
||||
0,
|
||||
embed_dim,
|
||||
2
|
||||
).float() *
|
||||
(-math.log(10000.0) / embed_dim)
|
||||
)
|
||||
|
||||
pe[:, 0::2] = torch.sin(position * div_term)
|
||||
pe[:, 1::2] = torch.cos(position * div_term)
|
||||
|
||||
pe = pe.unsqueeze(0)
|
||||
|
||||
# Stored as a buffer so it moves with the model
|
||||
# but isn't trained.
|
||||
self.register_buffer("pe", pe)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Args:
|
||||
x: Tensor of shape
|
||||
(batch_size, sequence_length, embed_dim)
|
||||
"""
|
||||
|
||||
seq_len = x.size(1)
|
||||
|
||||
x = x + self.pe[:, :seq_len]
|
||||
|
||||
return self.dropout(x)
|
||||
@@ -0,0 +1,196 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from config import (
|
||||
EMBED_DIM,
|
||||
NUM_HEADS,
|
||||
NUM_ENCODER_LAYERS,
|
||||
NUM_DECODER_LAYERS,
|
||||
FEED_FORWARD_DIM,
|
||||
DROPOUT,
|
||||
)
|
||||
|
||||
from model.positional_encoding import PositionalEncoding
|
||||
|
||||
|
||||
class MinecraftTransformer(nn.Module):
|
||||
"""
|
||||
Encoder-Decoder Transformer used by the Minecraft Builder AI.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
text_vocab_size,
|
||||
output_vocab_size,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.embed_dim = EMBED_DIM
|
||||
|
||||
# ----------------------------
|
||||
# Embeddings
|
||||
# ----------------------------
|
||||
|
||||
self.text_embedding = nn.Embedding(
|
||||
text_vocab_size,
|
||||
EMBED_DIM
|
||||
)
|
||||
|
||||
self.output_embedding = nn.Embedding(
|
||||
output_vocab_size,
|
||||
EMBED_DIM
|
||||
)
|
||||
|
||||
# ----------------------------
|
||||
# Positional Encoding
|
||||
# ----------------------------
|
||||
|
||||
self.text_position = PositionalEncoding(
|
||||
EMBED_DIM,
|
||||
DROPOUT
|
||||
)
|
||||
|
||||
self.output_position = PositionalEncoding(
|
||||
EMBED_DIM,
|
||||
DROPOUT
|
||||
)
|
||||
|
||||
# ----------------------------
|
||||
# Transformer
|
||||
# ----------------------------
|
||||
|
||||
self.transformer = nn.Transformer(
|
||||
d_model=EMBED_DIM,
|
||||
nhead=NUM_HEADS,
|
||||
num_encoder_layers=NUM_ENCODER_LAYERS,
|
||||
num_decoder_layers=NUM_DECODER_LAYERS,
|
||||
dim_feedforward=FEED_FORWARD_DIM,
|
||||
dropout=DROPOUT,
|
||||
batch_first=True,
|
||||
)
|
||||
|
||||
# ----------------------------
|
||||
# Output layer
|
||||
# ----------------------------
|
||||
|
||||
self.fc_out = nn.Linear(
|
||||
EMBED_DIM,
|
||||
output_vocab_size
|
||||
)
|
||||
|
||||
# ==================================================
|
||||
# Masks
|
||||
# ==================================================
|
||||
|
||||
def generate_square_subsequent_mask(self, size, device):
|
||||
"""
|
||||
Prevent the decoder from seeing future tokens.
|
||||
"""
|
||||
|
||||
return torch.triu(
|
||||
torch.full(
|
||||
(size, size),
|
||||
float("-inf"),
|
||||
device=device
|
||||
),
|
||||
diagonal=1
|
||||
)
|
||||
|
||||
# ==================================================
|
||||
# Forward
|
||||
# ==================================================
|
||||
|
||||
def forward(
|
||||
self,
|
||||
src,
|
||||
tgt,
|
||||
src_padding_mask=None,
|
||||
tgt_padding_mask=None,
|
||||
):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
src : (batch, src_len)
|
||||
|
||||
tgt : (batch, tgt_len)
|
||||
|
||||
Returns
|
||||
-------
|
||||
logits : (batch, tgt_len, output_vocab_size)
|
||||
"""
|
||||
|
||||
src = self.text_embedding(src)
|
||||
tgt = self.output_embedding(tgt)
|
||||
|
||||
src = self.text_position(src)
|
||||
tgt = self.output_position(tgt)
|
||||
|
||||
tgt_mask = self.generate_square_subsequent_mask(
|
||||
tgt.size(1),
|
||||
tgt.device
|
||||
)
|
||||
|
||||
output = self.transformer(
|
||||
src=src,
|
||||
tgt=tgt,
|
||||
tgt_mask=tgt_mask,
|
||||
src_key_padding_mask=src_padding_mask,
|
||||
tgt_key_padding_mask=tgt_padding_mask,
|
||||
memory_key_padding_mask=src_padding_mask,
|
||||
)
|
||||
|
||||
logits = self.fc_out(output)
|
||||
|
||||
return logits
|
||||
|
||||
# ==================================================
|
||||
# Encoder
|
||||
# ==================================================
|
||||
|
||||
def encode(
|
||||
self,
|
||||
src,
|
||||
src_padding_mask=None,
|
||||
):
|
||||
|
||||
src = self.text_embedding(src)
|
||||
src = self.text_position(src)
|
||||
|
||||
memory = self.transformer.encoder(
|
||||
src,
|
||||
src_key_padding_mask=src_padding_mask
|
||||
)
|
||||
|
||||
return memory
|
||||
|
||||
# ==================================================
|
||||
# Decoder
|
||||
# ==================================================
|
||||
|
||||
def decode(
|
||||
self,
|
||||
tgt,
|
||||
memory,
|
||||
tgt_padding_mask=None,
|
||||
memory_padding_mask=None,
|
||||
):
|
||||
|
||||
tgt = self.output_embedding(tgt)
|
||||
tgt = self.output_position(tgt)
|
||||
|
||||
tgt_mask = self.generate_square_subsequent_mask(
|
||||
tgt.size(1),
|
||||
tgt.device
|
||||
)
|
||||
|
||||
output = self.transformer.decoder(
|
||||
tgt=tgt,
|
||||
memory=memory,
|
||||
tgt_mask=tgt_mask,
|
||||
tgt_key_padding_mask=tgt_padding_mask,
|
||||
memory_key_padding_mask=memory_padding_mask,
|
||||
)
|
||||
|
||||
logits = self.fc_out(output)
|
||||
|
||||
return logits
|
||||
@@ -0,0 +1,5 @@
|
||||
EMBED_DIM = 512
|
||||
NUM_HEADS = 8
|
||||
NUM_ENCODER_LAYERS = 6
|
||||
NUM_DECODER_LAYERS = 6
|
||||
FEED_FORWARD_DIM = 2048
|
||||
@@ -0,0 +1,13 @@
|
||||
[project]
|
||||
name = "minebuildai"
|
||||
version = "0.1.0"
|
||||
requires-python = ">=3.12"
|
||||
dependencies = [
|
||||
"fastapi>=0.141.1",
|
||||
"mcschematic>=11.4.4",
|
||||
"noise>=1.2.2",
|
||||
"numpy>=2.5.1",
|
||||
"pydantic>=2.13.4",
|
||||
"requests>=2.34.2",
|
||||
"torch>=2.13.0",
|
||||
]
|
||||
@@ -0,0 +1,21 @@
|
||||
# Numerical computing
|
||||
numpy>=2.3.0
|
||||
|
||||
# Progress bars
|
||||
tqdm>=4.67.0
|
||||
|
||||
# Visualization (optional, useful during training)
|
||||
matplotlib>=3.10.0
|
||||
|
||||
# TensorBoard logging (optional)
|
||||
tensorboard>=2.20.0
|
||||
|
||||
# API Server
|
||||
fastapi>=0.139.0
|
||||
uvicorn>=0.35.0
|
||||
|
||||
# Data validation (FastAPI dependency, included automatically but pinned here)
|
||||
pydantic>=2.12.0
|
||||
|
||||
# MCSchematic
|
||||
mcschematic>=11.4.4
|
||||
@@ -0,0 +1,473 @@
|
||||
import json
|
||||
|
||||
import requests
|
||||
from fastapi import FastAPI
|
||||
from pydantic import BaseModel
|
||||
|
||||
from config import DEVICE, CHECKPOINT_DIR
|
||||
from data.tokenizer import Tokenizer
|
||||
from makescript import json_string_to_py
|
||||
from model.builder_model import BuilderModel
|
||||
from utils.checkpoint import latest_checkpoint, load_checkpoint
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
import os
|
||||
import base64
|
||||
|
||||
app = FastAPI(title="Minecraft Builder AI")
|
||||
|
||||
if DEVICE == "cuda":
|
||||
print("Cuda Mode: true")
|
||||
else:
|
||||
print("Cuda Mode: false")
|
||||
|
||||
print("Loading tokenizer...")
|
||||
tokenizer = Tokenizer()
|
||||
tokenizer.load("data")
|
||||
|
||||
print("Loading model...")
|
||||
model = BuilderModel(tokenizer).to(DEVICE)
|
||||
|
||||
checkpoint = latest_checkpoint(CHECKPOINT_DIR)
|
||||
load_checkpoint(
|
||||
checkpoint,
|
||||
model,
|
||||
optimizer=None,
|
||||
map_location=DEVICE,
|
||||
)
|
||||
|
||||
model.eval()
|
||||
|
||||
|
||||
class BuildRequest(BaseModel):
|
||||
prompt: str
|
||||
|
||||
|
||||
@app.post("/generate")
|
||||
def generate(req: BuildRequest):
|
||||
|
||||
blocks, tokens = model.generate(req.prompt)
|
||||
|
||||
return {
|
||||
"prompt": req.prompt,
|
||||
"tokens": tokens,
|
||||
"blocks": blocks,
|
||||
}
|
||||
|
||||
class RunRequest(BaseModel):
|
||||
script: str
|
||||
|
||||
@app.post("/run_script")
|
||||
def runscript(req: RunRequest):
|
||||
result = subprocess.run([
|
||||
sys.executable,
|
||||
"builders/" + req.script + ".py"
|
||||
])
|
||||
|
||||
if result.returncode != 0:
|
||||
return {
|
||||
"success": False,
|
||||
"schem": None
|
||||
}
|
||||
|
||||
schem_path = "schematics/" + req.script + ".schem"
|
||||
|
||||
if not os.path.exists(schem_path):
|
||||
return {
|
||||
"success": False,
|
||||
"schem": None
|
||||
}
|
||||
|
||||
with open(schem_path, "rb") as f:
|
||||
schem = base64.b64encode(f.read()).decode("ascii")
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"schem": schem
|
||||
}
|
||||
|
||||
class AIScriptRequest(BaseModel):
|
||||
url: str
|
||||
model: str
|
||||
api_key: str
|
||||
prompt: str
|
||||
context_length: int
|
||||
temperature: float
|
||||
reasoning: str
|
||||
|
||||
@app.post("/aiscript")
|
||||
def aiscript(req: AIScriptRequest):
|
||||
systemprompt = r"""
|
||||
You are a Minecraft Building AI.
|
||||
You build with the mcschematic python library.
|
||||
Your task is to return a valid script in the following format:
|
||||
{
|
||||
"file_name": FILE_NAME,
|
||||
"script: SCRIPT,
|
||||
"response": RESPONSE
|
||||
}
|
||||
|
||||
Rules:
|
||||
- Only use valid Minecraft blocks
|
||||
- The output folder is called "schematics"
|
||||
- The schematic name must match with the FILE_NAME
|
||||
- In settings there always needs to be a variable called "SEED"
|
||||
|
||||
Documentation:
|
||||
mcschematic is a Python library for programmatically creating Minecraft Java Edition schematic files. The basic workflow is:
|
||||
Create an MCSchematic object.
|
||||
Place Minecraft blocks at (x, y, z) coordinates with setBlock().
|
||||
Optionally manipulate blocks/entities/data.
|
||||
Save the resulting structure as a .schem file with save().
|
||||
The core conceptual model is:
|
||||
import mcschematic
|
||||
schematic = mcschematic.MCSchematic()
|
||||
schematic.setBlock((x, y, z), "minecraft:oak_log")
|
||||
schematic.setBlock((x, y, z), "minecraft:oak_leaves")
|
||||
schematic.save(
|
||||
"output_directory",
|
||||
"my_structure",
|
||||
mcschematic.Version.JE_1_20_1
|
||||
)
|
||||
Important concepts
|
||||
MCSchematic
|
||||
This is the main container representing the Minecraft structure.
|
||||
schem = mcschematic.MCSchematic()
|
||||
Think of it as an in-memory 3D voxel structure.
|
||||
Coordinates
|
||||
Blocks are addressed using:
|
||||
(x, y, z)
|
||||
For procedural generation, treat these as a normal 3D Cartesian coordinate system. Pick a clear origin convention—for example, (0, 0, 0) as the bottom-center of a tree—and build everything relative to that.
|
||||
Placing blocks
|
||||
The primary operation is:
|
||||
schem.setBlock((x, y, z), "minecraft:oak_log")
|
||||
Block IDs should use Minecraft's namespaced IDs:
|
||||
minecraft:stone
|
||||
minecraft:dirt
|
||||
minecraft:oak_log
|
||||
minecraft:oak_leaves
|
||||
minecraft:glass
|
||||
Block states can be specified as part of the block string when supported, e.g.:
|
||||
"minecraft:oak_leaves[persistent=true]"
|
||||
For a procedural generator, create helper functions around setBlock() rather than scattering raw calls throughout the code.
|
||||
Example:
|
||||
def block(schem, x, y, z, block_id):
|
||||
schem.setBlock((x, y, z), block_id)
|
||||
Then your generation code can operate at a higher level:
|
||||
block(schem, x, y, z, "minecraft:oak_log")
|
||||
Procedural generation
|
||||
mcschematic isn't a tree/terrain generator itself. It provides the Minecraft schematic representation and file output.
|
||||
The procedural geometry should be implemented by the agent.
|
||||
For example:
|
||||
Tree generator
|
||||
│
|
||||
├── trunk()
|
||||
│ └── places log blocks
|
||||
│
|
||||
├── branch()
|
||||
│ └── calculates a 3D line and places logs
|
||||
│
|
||||
├── roots()
|
||||
│ └── creates radial structures
|
||||
│
|
||||
└── canopy()
|
||||
└── creates randomized leaf volumes
|
||||
│
|
||||
└── setBlock()
|
||||
Useful geometric primitives to implement yourself include:
|
||||
line between two 3D points
|
||||
cylinder
|
||||
sphere / ellipsoid
|
||||
cone
|
||||
tapered branch
|
||||
radial distribution
|
||||
noise/randomized volume
|
||||
distance-based falloff
|
||||
For example, a spherical leaf cluster can be generated by checking:
|
||||
distance = (
|
||||
(x - cx) ** 2 +
|
||||
(y - cy) ** 2 +
|
||||
(z - cz) ** 2
|
||||
)
|
||||
if distance <= radius ** 2:
|
||||
schem.setBlock(
|
||||
(x, y, z),
|
||||
"minecraft:oak_leaves[persistent=true]"
|
||||
)
|
||||
Minecraft version
|
||||
The library has Minecraft-version constants under mcschematic.Version. The agent should inspect the installed library rather than hard-code a version based on memory.
|
||||
For example:
|
||||
print(dir(mcschematic.Version))
|
||||
and choose the appropriate Java Edition version available in the installed package.
|
||||
Likewise, the agent should inspect the actual installed API:
|
||||
import mcschematic
|
||||
print(mcschematic.version)
|
||||
print(dir(mcschematic))
|
||||
If there is uncertainty about an API call, inspect the package source/signatures instead of inventing an API.
|
||||
Saving
|
||||
The schematic is eventually serialized with:
|
||||
schem.save(
|
||||
output_folder,
|
||||
schematic_name,
|
||||
version
|
||||
)
|
||||
The resulting file is intended to be usable by Minecraft schematic-compatible tools.
|
||||
Recommended architecture for a tree generator
|
||||
A good design would be:
|
||||
tree = TreeGenerator(
|
||||
height=20,
|
||||
trunk_radius=2,
|
||||
seed=1234,
|
||||
)
|
||||
tree.generate()
|
||||
tree.export(
|
||||
"output/tree.schem"
|
||||
)
|
||||
Internally:
|
||||
TreeGenerator
|
||||
↓
|
||||
generate_trunk()
|
||||
↓
|
||||
generate_branches()
|
||||
↓
|
||||
generate_roots()
|
||||
↓
|
||||
generate_canopy()
|
||||
↓
|
||||
MCSchematic.setBlock()
|
||||
↓
|
||||
MCSchematic.save()
|
||||
|
||||
Helper Methods:
|
||||
def set_block(x, y, z, block):
|
||||
schem.setBlock((x, y, z), block)
|
||||
|
||||
def fill_sphere(cx, cy, cz, radius, block, irregularity=0.0):
|
||||
# Create a branch between two points.
|
||||
for x in range(cx - radius, cx + radius + 1):
|
||||
for y in range(cy - radius, cy + radius + 1):
|
||||
for z in range(cz - radius, cz + radius + 1):
|
||||
|
||||
distance = math.sqrt(
|
||||
(x - cx) ** 2 +
|
||||
(y - cy) ** 2 +
|
||||
(z - cz) ** 2
|
||||
)
|
||||
|
||||
# Make the sphere slightly irregular.
|
||||
wobble = random.uniform(-irregularity, irregularity)
|
||||
|
||||
if distance <= radius + wobble:
|
||||
set_block(x, y, z, block)
|
||||
|
||||
|
||||
def cylinder(cx, cz, y_start, y_end, radius, block):
|
||||
# Create a vertical cylinder.
|
||||
for y in range(y_start, y_end + 1):
|
||||
for x in range(cx - radius, cx + radius + 1):
|
||||
for z in range(cz - radius, cz + radius + 1):
|
||||
|
||||
if (x - cx) ** 2 + (z - cz) ** 2 <= radius ** 2:
|
||||
set_block(x, y, z, block)
|
||||
|
||||
|
||||
def branch(x1, y1, z1, x2, y2, z2, radius=1):
|
||||
# Create a branch between two points.
|
||||
steps = max(
|
||||
abs(x2 - x1),
|
||||
abs(y2 - y1),
|
||||
abs(z2 - z1)
|
||||
)
|
||||
|
||||
for i in range(steps + 1):
|
||||
t = i / max(steps, 1)
|
||||
|
||||
x = round(x1 + (x2 - x1) * t)
|
||||
y = round(y1 + (y2 - y1) * t)
|
||||
z = round(z1 + (z2 - z1) * t)
|
||||
|
||||
cylinder(x, z, y, y, radius, "minecraft:oak_log")
|
||||
|
||||
These are example helper methods. Also try to make your own when needed.
|
||||
|
||||
This is how the schematic will be saved:
|
||||
schem.save(
|
||||
OUTPUT_FOLDER,
|
||||
OUTPUT_NAME,
|
||||
MC_VERSION
|
||||
)
|
||||
|
||||
print(f"Saved {OUTPUT_FOLDER}/{OUTPUT_NAME}.schem")
|
||||
|
||||
Coding Rules ONLY APPLY THESE IN THE SCRIPT NOT THE ENTIRE JSON:
|
||||
file_name:
|
||||
Use only letters, numbers, underscores, and hyphens.
|
||||
Do not include .py.
|
||||
Do not include directories or paths.
|
||||
script:
|
||||
Must contain the complete, executable Python source code.
|
||||
The script must use mcschematic.
|
||||
The script must save the generated schematic using schem.save(...).
|
||||
Use Minecraft Java Edition 1.20.1 unless the user specifies another version.
|
||||
The script must be self-contained and must not require the user to manually modify it.
|
||||
CRITICAL JSON REQUIREMENT:
|
||||
The value of script is a JSON string.
|
||||
You MUST properly JSON-escape the Python source code.
|
||||
Every double quote (") INSIDE the Python script must be escaped as \".
|
||||
Every newline inside the Python script must be represented as \n.
|
||||
Never output literal unescaped newlines inside the JSON script string.
|
||||
Never produce invalid JSON.
|
||||
Do NOT wrap the response in Markdown code fences.
|
||||
Do NOT output json. Do NOT output python.
|
||||
Output ONLY the JSON object.
|
||||
Do not add explanations, comments outside the Python script, or any text before or after the JSON object.
|
||||
Before responding, internally verify that your entire response is valid JSON and could be successfully parsed with Python's json.loads().
|
||||
The resulting value of:
|
||||
json.loads(response)["script"]
|
||||
must be directly writable to a .py file and must produce valid Python source code.
|
||||
In "response" give a short description what you build.
|
||||
Example of the REQUIRED output format:
|
||||
{
|
||||
"file_name": "small_oak_tree",
|
||||
"script": "import mcschematic\nimport math\n\nMC_VERSION = mcschematic.Version.JE_1_20_1\nschem = mcschematic.MCSchematic()\n\nschem.setBlock((0, 0, 0), \"minecraft\")\n\nschem.save(\"schematics\", \"small_oak_tree\", MC_VERSION)\n",
|
||||
"response": "I build your small oak tree."
|
||||
}
|
||||
Remember: the example above demonstrates the JSON encoding requirements. Your actual response must contain ONLY one valid JSON object.
|
||||
"""
|
||||
|
||||
payload = {
|
||||
"model": req.model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": systemprompt
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": req.prompt
|
||||
}
|
||||
],
|
||||
"max_completion_tokens": req.context_length,
|
||||
"reasoning_effort": req.reasoning,
|
||||
"temperature": req.temperature,
|
||||
"stream": True,
|
||||
}
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {req.api_key}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
req.url + "/chat/completions",
|
||||
json=payload,
|
||||
headers=headers,
|
||||
timeout=600,
|
||||
stream=True
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
reasoning = ""
|
||||
answer = ""
|
||||
|
||||
for line in response.iter_lines(decode_unicode=True):
|
||||
if not line or not line.startswith("data: "):
|
||||
continue
|
||||
|
||||
data = line[6:]
|
||||
|
||||
if data == "[DONE]":
|
||||
break
|
||||
|
||||
try:
|
||||
chunk = json.loads(data)
|
||||
except json.JSONDecodeError:
|
||||
print("Could not decode SSE chunk:", data)
|
||||
continue
|
||||
|
||||
# Some chunks contain no choices.
|
||||
choices = chunk.get("choices", [])
|
||||
if not choices:
|
||||
continue
|
||||
|
||||
choice = choices[0]
|
||||
delta = choice.get("delta", {})
|
||||
|
||||
# Reasoning tokens
|
||||
reasoning_token = delta.get("reasoning_content")
|
||||
if reasoning_token:
|
||||
reasoning += reasoning_token
|
||||
print(reasoning_token, end="", flush=True)
|
||||
|
||||
# Normal answer tokens
|
||||
content_token = delta.get("content")
|
||||
if content_token:
|
||||
answer += content_token
|
||||
print(content_token, end="", flush=True)
|
||||
|
||||
print("\n\nFINAL ANSWER:")
|
||||
print(answer)
|
||||
|
||||
# Parse the AI's JSON
|
||||
try:
|
||||
result = json.loads(answer)
|
||||
except json.JSONDecodeError as e:
|
||||
print("AI did not return valid JSON!")
|
||||
print("JSON error:", e)
|
||||
print("Raw answer:")
|
||||
print(repr(answer))
|
||||
|
||||
return {
|
||||
"success": False,
|
||||
"file_name": "failed",
|
||||
"response": "AI returned invalid JSON"
|
||||
}
|
||||
|
||||
# Validate expected fields
|
||||
if "file_name" not in result:
|
||||
raise ValueError("AI JSON is missing 'file_name'")
|
||||
|
||||
if "script" not in result:
|
||||
raise ValueError("AI JSON is missing 'script'")
|
||||
|
||||
if "response" not in result:
|
||||
raise ValueError("AI JSON is missing 'response'")
|
||||
|
||||
file_name = result["file_name"]
|
||||
script = result["script"]
|
||||
remessage = result["response"]
|
||||
|
||||
# Write the actual Python script
|
||||
with open(file_name + ".py", "w", encoding="utf-8") as f:
|
||||
f.write(script)
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"file_name": file_name,
|
||||
"response": remessage
|
||||
}
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
print("API request failed!")
|
||||
print("Error:", e)
|
||||
|
||||
if e.response is not None:
|
||||
print("Status code:", e.response.status_code)
|
||||
print("Response:", e.response.text)
|
||||
|
||||
return {
|
||||
"success": False,
|
||||
"file_name": "failed",
|
||||
"response": "failed"
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
print("Unexpected error:", e)
|
||||
|
||||
return {
|
||||
"success": False,
|
||||
"file_name": "failed",
|
||||
"response": str(e)
|
||||
}
|
||||
@@ -0,0 +1,156 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
|
||||
from config import (
|
||||
DATASET_PATH,
|
||||
BATCH_SIZE,
|
||||
LEARNING_RATE,
|
||||
EPOCHS,
|
||||
DEVICE,
|
||||
CHECKPOINT_DIR,
|
||||
)
|
||||
|
||||
from data.tokenizer import Tokenizer
|
||||
from data.dataset import create_dataloader
|
||||
|
||||
from model.builder_model import BuilderModel
|
||||
|
||||
from utils.checkpoint import (
|
||||
save_checkpoint,
|
||||
latest_checkpoint,
|
||||
load_checkpoint,
|
||||
)
|
||||
|
||||
def format_params(n):
|
||||
if n >= 1_000_000_000:
|
||||
return f"{n:,} ({n/1_000_000_000:.2f}B)"
|
||||
elif n >= 1_000_000:
|
||||
return f"{n:,} ({n/1_000_000:.2f}M)"
|
||||
elif n >= 1_000:
|
||||
return f"{n:,} ({n/1_000:.2f}K)"
|
||||
return str(n)
|
||||
|
||||
def train():
|
||||
|
||||
print("=" * 60)
|
||||
print("Building tokenizer...")
|
||||
print("=" * 60)
|
||||
|
||||
tokenizer = Tokenizer()
|
||||
tokenizer.build(DATASET_PATH)
|
||||
tokenizer.save("data")
|
||||
|
||||
print("Prompt vocabulary :", tokenizer.text_vocab_size)
|
||||
print("Output vocabulary :", tokenizer.output_vocab_size)
|
||||
|
||||
print("=" * 60)
|
||||
print("Loading dataset...")
|
||||
print("=" * 60)
|
||||
|
||||
dataloader = create_dataloader(
|
||||
DATASET_PATH,
|
||||
tokenizer,
|
||||
BATCH_SIZE,
|
||||
)
|
||||
|
||||
print("Creating model...")
|
||||
|
||||
model = BuilderModel(tokenizer).to(DEVICE)
|
||||
|
||||
optimizer = optim.AdamW(
|
||||
model.parameters(),
|
||||
lr=LEARNING_RATE,
|
||||
)
|
||||
|
||||
criterion = nn.CrossEntropyLoss(
|
||||
ignore_index=tokenizer.output_to_id["<PAD>"]
|
||||
)
|
||||
|
||||
start_epoch = 1
|
||||
|
||||
checkpoint = latest_checkpoint(CHECKPOINT_DIR)
|
||||
|
||||
if checkpoint:
|
||||
|
||||
print("Loading checkpoint:", checkpoint)
|
||||
|
||||
start_epoch, _ = load_checkpoint(
|
||||
checkpoint,
|
||||
model,
|
||||
optimizer,
|
||||
DEVICE,
|
||||
)
|
||||
|
||||
start_epoch += 1
|
||||
|
||||
print("=" * 60)
|
||||
print("Training")
|
||||
print("=" * 60)
|
||||
|
||||
for epoch in range(start_epoch, EPOCHS + 1):
|
||||
|
||||
model.train()
|
||||
|
||||
total_loss = 0
|
||||
|
||||
for prompts, outputs in dataloader:
|
||||
|
||||
prompts = prompts.to(DEVICE)
|
||||
outputs = outputs.to(DEVICE)
|
||||
|
||||
decoder_input = outputs[:, :-1]
|
||||
targets = outputs[:, 1:]
|
||||
|
||||
optimizer.zero_grad()
|
||||
|
||||
logits = model(
|
||||
prompts,
|
||||
decoder_input,
|
||||
)
|
||||
|
||||
loss = criterion(
|
||||
logits.reshape(-1, logits.size(-1)),
|
||||
targets.reshape(-1),
|
||||
)
|
||||
|
||||
loss.backward()
|
||||
|
||||
torch.nn.utils.clip_grad_norm_(
|
||||
model.parameters(),
|
||||
1.0,
|
||||
)
|
||||
|
||||
optimizer.step()
|
||||
|
||||
total_loss += loss.item()
|
||||
|
||||
avg_loss = total_loss / len(dataloader)
|
||||
|
||||
print(
|
||||
f"Epoch {epoch}/{EPOCHS} | Loss: {avg_loss:.4f}"
|
||||
)
|
||||
|
||||
print("\nTraining finished!")
|
||||
|
||||
total_params = sum(p.numel() for p in model.parameters())
|
||||
trainable_params = sum(
|
||||
p.numel() for p in model.parameters() if p.requires_grad
|
||||
)
|
||||
|
||||
print("=" * 60)
|
||||
print(f"Total parameters: {format_params(total_params)}")
|
||||
print(f"Trainable parameters: {format_params(trainable_params)}")
|
||||
print("=" * 60)
|
||||
|
||||
save_checkpoint(
|
||||
model,
|
||||
optimizer,
|
||||
EPOCHS,
|
||||
avg_loss,
|
||||
CHECKPOINT_DIR / "final_model.pth",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
train()
|
||||
@@ -0,0 +1,85 @@
|
||||
import os
|
||||
import torch
|
||||
|
||||
|
||||
def save_checkpoint(
|
||||
model,
|
||||
optimizer,
|
||||
epoch,
|
||||
loss,
|
||||
path,
|
||||
):
|
||||
"""
|
||||
Save a training checkpoint.
|
||||
"""
|
||||
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
|
||||
torch.save(
|
||||
{
|
||||
"epoch": epoch,
|
||||
"model_state_dict": model.state_dict(),
|
||||
"optimizer_state_dict": optimizer.state_dict(),
|
||||
"loss": loss,
|
||||
},
|
||||
path,
|
||||
)
|
||||
|
||||
|
||||
def load_checkpoint(
|
||||
path,
|
||||
model,
|
||||
optimizer=None,
|
||||
map_location="cpu",
|
||||
):
|
||||
"""
|
||||
Load a training checkpoint.
|
||||
|
||||
Returns:
|
||||
epoch (int)
|
||||
loss (float)
|
||||
"""
|
||||
|
||||
checkpoint = torch.load(
|
||||
path,
|
||||
map_location=map_location,
|
||||
)
|
||||
|
||||
model.load_state_dict(
|
||||
checkpoint["model_state_dict"]
|
||||
)
|
||||
|
||||
if optimizer is not None:
|
||||
optimizer.load_state_dict(
|
||||
checkpoint["optimizer_state_dict"]
|
||||
)
|
||||
|
||||
return (
|
||||
checkpoint["epoch"],
|
||||
checkpoint["loss"],
|
||||
)
|
||||
|
||||
|
||||
def latest_checkpoint(folder):
|
||||
"""
|
||||
Returns the newest checkpoint file in a folder.
|
||||
"""
|
||||
|
||||
if not os.path.exists(folder):
|
||||
return None
|
||||
|
||||
checkpoints = [
|
||||
os.path.join(folder, file)
|
||||
for file in os.listdir(folder)
|
||||
if file.endswith(".pth")
|
||||
]
|
||||
|
||||
if len(checkpoints) == 0:
|
||||
return None
|
||||
|
||||
checkpoints.sort(
|
||||
key=os.path.getmtime,
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
return checkpoints[0]
|
||||
Reference in New Issue
Block a user