import json
import torch

from huggingface_hub import hf_hub_download

# Download model files

model_path = hf_hub_download(
    repo_id="tenperformer/dungeon-master-gpt",
    filename="best_dungeonmastergpt.pt"
)

config_path = hf_hub_download(
    repo_id="tenperformer/dungeon-master-gpt",
    filename="gpt_config.json"
)

# Load configuration

with open(config_path, "r") as f:
    cfg = json.load(f)

config = GPTConfig(
    vocab_size=cfg["vocab_size"],
    block_size=cfg["block_size"],
    n_layer=cfg["n_layer"],
    n_head=cfg["n_head"],
    n_embd=cfg["n_embd"],
    dropout=cfg["dropout"]
)

# Create model

device = "cuda" if torch.cuda.is_available() else "cpu"

model = GPT(config).to(device)

# Load weights

model.load_state_dict(
    torch.load(
        model_path,
        map_location=device
    )
)

model.eval()

# Generate text

prompt = "The wizard entered the ancient dungeon"

context = torch.tensor(
    enc.encode(prompt),
    dtype=torch.long
).unsqueeze(0).to(device)

output = model.generate(
    context,
    max_new_tokens=200
)

print(
    enc.decode(output[0].tolist())
)