74 lines
1.5 KiB
Python
74 lines
1.5 KiB
Python
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() |