init commit

.gitignore

@@ -0,0 +1,54 @@
+*.py[cod]
+
+# C extensions
+*.so
+
+# Packages
+*.egg
+*.egg-info
+dist
+build
+eggs
+parts
+bin
+var
+sdist
+develop-eggs
+.installed.cfg
+lib
+lib64
+__pycache__
+
+# Installer logs
+pip-log.txt
+
+# Unit test / coverage reports
+.coverage
+.tox
+nosetests.xml
+
+# Translations
+*.mo
+
+# Mr Developer
+.mr.developer.cfg
+.project
+.pydevproject
+test.json
+*.pickle
+venv
+.idea
+*.vscode/
+.DS_Store
+
+#notebooks
+*/**/.ipynb_checkpoints/*
+
+# logs
+*/**/checkpoints/*
+*/**/mlruns/*
+*/**/tensorboard_logs/*
+*/**/wandb/*
+checkpoints/*
+wandb/*
+mlruns/*

LICENSE

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2023 Shyam Sudhakaran
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

app.py

@@ -0,0 +1,446 @@
+import gradio as gr
+import uuid
+from mario_gpt.lm import MarioLM
+from mario_gpt.utils import convert_level_to_png
+
+from fastapi import FastAPI
+from fastapi.staticfiles import StaticFiles
+
+import uvicorn
+import boto3
+import json
+
+bedrock_runtime = boto3.client(
+    service_name="bedrock-runtime",
+    region_name="us-east-1",
+)
+
+
+def get_raw_text(level_data):
+    raw_text = ""
+    for line in level_data:
+        raw_text += line + "\n"
+    return raw_text
+
+
+def combine_levels(level_arrays):
+    num_rows = len(level_arrays[0])
+
+    combined_level = []
+
+    for row in range(num_rows):
+        combined_row = ""
+        for level in level_arrays:
+            combined_row += level[row]
+        combined_level.append(combined_row)
+
+    return combined_level
+
+
+def write_level_to_file(level_data, file_name):
+    with open(file_name, "w") as file:
+        for line in level_data:
+            file.write(line + "\n")
+
+
+def clean_level_data(input_string):
+    # Find the start and end indices of the level data
+    start_index = input_string.find("[")
+    end_index = input_string.rfind("]")
+
+    # Extract the level data
+    level_data = input_string[start_index + 1 : end_index]
+
+    # Split the level data into lines
+    lines = level_data.split(",")
+
+    # Clean each line
+    cleaned_lines = []
+    for line in lines:
+        # Remove leading and trailing whitespace and quotes
+        cleaned_line = line.strip().strip("'")
+
+        # Ensure the line has exactly 50 characters
+        if len(cleaned_line) < 50:
+            cleaned_line += "-" * (50 - len(cleaned_line))
+        elif len(cleaned_line) > 50:
+            cleaned_line = cleaned_line[:50]
+
+        cleaned_lines.append(cleaned_line)
+
+    return cleaned_lines
+
+
+def call_llama3_70b(system_prompt, prompt):
+    prompt_config = {
+        "prompt": system_prompt + prompt,
+        "max_gen_len": 2048,
+        "top_p": 0.9,
+        "temperature": 0.7,
+    }
+
+    body = json.dumps(prompt_config)
+
+    modelId = "meta.llama3-70b-instruct-v1:0"
+    accept = "application/json"
+    contentType = "application/json"
+
+    response = bedrock_runtime.invoke_model(
+        body=body, modelId=modelId, accept=accept, contentType=contentType
+    )
+    response_body = json.loads(response.get("body").read())
+
+    results = response_body["generation"].strip()
+    return results
+
+
+def call_llama3_8b(system_prompt, prompt):
+    prompt_config = {
+        "prompt": system_prompt + prompt,
+        "max_gen_len": 2048,
+        "top_p": 0.9,
+        "temperature": 0.7,
+    }
+
+    body = json.dumps(prompt_config)
+
+    modelId = "meta.llama3-8b-instruct-v1:0"
+    accept = "application/json"
+    contentType = "application/json"
+
+    response = bedrock_runtime.invoke_model(
+        body=body, modelId=modelId, accept=accept, contentType=contentType
+    )
+    response_body = json.loads(response.get("body").read())
+
+    results = response_body["generation"].strip()
+    return results
+
+
+# def call_claude_3_opus(system_prompt, prompt):
+
+#     prompt_config = {
+#         "anthropic_version": "bedrock-2023-05-31",
+#         "max_tokens": 4096,
+#         "system": system_prompt,
+#         "messages": [
+#             {
+#                 "role": "user",
+#                 "content": [
+#                     {"type": "text", "text": prompt},
+#                 ],
+#             }
+#         ],
+#     }
+
+#     body = json.dumps(prompt_config)
+
+#     modelId = "anthropic.claude-3-opus-20240229-v1:0"
+#     accept = "application/json"
+#     contentType = "application/json"
+
+#     response = bedrock_runtime.invoke_model(
+#         body=body, modelId=modelId, accept=accept, contentType=contentType
+#     )
+#     response_body = json.loads(response.get("body").read())
+
+#     results = response_body.get("content")[0].get("text")
+#     return results
+
+
+# Call Claude model
+def call_claude_3_sonnet(system_prompt, prompt):
+
+    prompt_config = {
+        "anthropic_version": "bedrock-2023-05-31",
+        "max_tokens": 4096,
+        "system": system_prompt,
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt},
+                ],
+            }
+        ],
+    }
+
+    body = json.dumps(prompt_config)
+
+    modelId = "anthropic.claude-3-sonnet-20240229-v1:0"
+    accept = "application/json"
+    contentType = "application/json"
+
+    response = bedrock_runtime.invoke_model(
+        body=body, modelId=modelId, accept=accept, contentType=contentType
+    )
+    response_body = json.loads(response.get("body").read())
+
+    results = response_body.get("content")[0].get("text")
+    return results
+
+
+def call_claude_3_haiku(system_prompt, prompt):
+
+    prompt_config = {
+        "anthropic_version": "bedrock-2023-05-31",
+        "max_tokens": 4096,
+        "system": system_prompt,
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt},
+                ],
+            }
+        ],
+    }
+
+    body = json.dumps(prompt_config)
+
+    modelId = "anthropic.claude-3-haiku-20240307-v1:0"
+    accept = "application/json"
+    contentType = "application/json"
+
+    response = bedrock_runtime.invoke_model(
+        body=body, modelId=modelId, accept=accept, contentType=contentType
+    )
+    response_body = json.loads(response.get("body").read())
+
+    results = response_body.get("content")[0].get("text")
+    return results
+
+
+system_prompt_text = """
+As an esteemed level designer renowned for creating some of the top 100 levels in Super Mario Maker, you are tasked with crafting a playable section for the original Super Mario on NES. Your extensive experience and creativity are key to designing levels that are not only challenging but also immensely enjoyable. Use the following symbols to represent different game elements, ensuring each level is a masterpiece of design:
+
+<symbols>
+- = "Sky"
+X = "Unbreakable Block"
+E = "Enemy"
+o = "Coin"
+S = "Breakable Block"
+? = "Question Block"
+[] = "Pipe"
+<> = "End of Pipe"
+</symbols>
+
+
+Adhere to these level layout specifications:
+
+<level guidelines>
+Pipes should be vertical and follow this format:
+<>
+[]
+[]
+
+Ensure there is a clear and navigable path that Mario can follow from the start to the end of the level. This path may involve jumping on blocks or pipes, running on blocks. 
+
+The path should be continuous and not lead Mario into any dead ends or impossible situations.
+
+Place unbreakable blocks (X) or other platform elements strategically to create a solid foundation for Mario to walk on. Avoid creating large gaps or sections without any ground or platforms, as Mario needs a surface to stand on.
+
+Adjust the complexity and elements based on the specific level request, ensuring that Mario can always complete the level successfully by following the designated path.
+</level guidelines>
+
+
+
+For example a prompt that asks for a level with one pipe and blocks with 2 Goombas. 
+
+Here is an example output:
+<example>
+['--------------------------------------------------',
+'--------------------------------------------------',
+'--------------------------------------------------',
+'--------------------------------------------------',
+'-------------------------------------------------o',
+'--------XSSSSS---------------------------------SSS',
+'--------X-----------------------------------------',
+'--------X-----------------------------------------',
+'-------EX--E-X---------------xxxx-?-----------xxxx',
+'--------XSS?SX---QQ?QQ------xx<>-x-----------xx--?',
+'---------------------------xx-[]--x---------xx----',
+'--------------------------xx--[]---x-------xx-----',
+'xxxxxxxxxxxxxxxxxxxxxxxxxxx---[]----xxxxxxxx------',
+'XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX---XXX']
+</example>
+
+
+Generate the level section as a 2D array, where each row is represented as a string of characters. The level section should be 14 rows tall and 50 columns wide. Only return the 2D array of characters.
+
+Remember, your creations should challenge players but remain fair. Use your expertise to weave together obstacles and rewards, encouraging exploration and skillful play. Always ensure that Mario has a clear and navigable route to finish the level, and provide ample block tiles for Mario to walk on.
+"""
+
+
+mario_lm = MarioLM()
+# device = torch.device('cuda')
+# mario_lm = mario_lm.to(device)
+TILE_DIR = "mario_gpt/data/tiles"
+
+app = FastAPI()
+
+
+def make_html_file(generated_level):
+    level_text = generated_level
+    unique_id = uuid.uuid1()
+    with open(f"static/demo-{unique_id}.html", "w", encoding="utf-8") as f:
+        f.write(
+            f"""<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="utf-8">
+    <title>Mario Game</title>
+    <script src="https://cjrtnc.leaningtech.com/20230216/loader.js"></script>
+</head>
+<body>
+</body>
+<script>
+    cheerpjInit().then(function () {{
+        cheerpjAddStringFile("/str/mylevel.txt", `{level_text}`);
+    }});
+    cheerpjCreateDisplay(612, 600);
+    cheerpjRunJar("/app/static/mario.jar");
+</script>
+</html>"""
+        )
+    return f"demo-{unique_id}.html"
+
+
+def generate(model, prompt, system_prompt=system_prompt_text):
+
+    print(f"Using prompt: {prompt}")
+
+    if system_prompt == "":
+        system_prompt = system_prompt_text
+
+    # #  prompt 3 times
+    # prompts = [prompt, prompt, prompt]
+
+    # levels_array = []
+
+    # for index, prompt in enumerate(prompts):
+
+    #     level = call_claude_3_sonnet(system_prompt, prompt)
+    #     cleaned_level = clean_level_data(level)
+
+    #     levels_array.append(cleaned_level)
+
+    # final_level = combine_levels(levels_array)
+    # raw_level_text = get_raw_text(final_level)
+
+    if model == "Claude Sonnet":
+        level = call_claude_3_sonnet(system_prompt, prompt)
+    elif model == "Claude Haiku":
+        level = call_claude_3_haiku(system_prompt, prompt)
+    elif model == "Llama3 70B":
+        level = call_llama3_70b(system_prompt, prompt)
+    elif model == "Llama3 8B":
+        level = call_llama3_8b(system_prompt, prompt)
+    # elif model == "Cladue Opus":
+    #     level = call_claude_3_opus(system_prompt, prompt)
+    else:
+        raise ValueError("Invalid model")
+
+    # level = call_claude_3_sonnet(system_prompt, prompt)
+    cleaned_level = clean_level_data(level)
+    raw_level_text = get_raw_text(cleaned_level)
+
+    filename = make_html_file(raw_level_text)
+    img = convert_level_to_png(cleaned_level, mario_lm.tokenizer)[0]
+
+    gradio_html = f"""<div>
+        <iframe width=612 height=612 style="margin: 0 auto" src="static/{filename}"></iframe>
+        <p style="text-align:center">Press the arrow keys to move. Press <code>a</code> to run, <code>s</code> to jump and <code>d</code> to shoot fireflowers</p>
+    </div>"""
+    return [img, gradio_html]
+
+
+with gr.Blocks().queue() as demo:
+    gr.Markdown(
+        """### Playable demo for MarioGPT: Open-Ended Text2Level Generation through Large Language Models - Amazon Bedrock Edition
+    [[Github](https://github.com/shyamsn97/mario-gpt)], [[Paper](https://arxiv.org/abs/2302.05981), [[Amazon Bedrock](https://docs.aws.amazon.com/bedrock/latest/userguide/service_code_examples.html?trk=2403b700-9ee9-49e8-aed8-411dea5cf5ae&sc_channel=el)]
+    """
+    )
+    with gr.Tabs():
+        with gr.TabItem("Prompt Settings"):
+
+            with gr.Accordion(label="System Prompt", open=False):
+                #     temperature = gr.Number(
+                #         value=2.0,
+                #         label="temperature: Increase these for more diverse, but lower quality, generations",
+                #     )
+                system_prompt = gr.TextArea(
+                    value=system_prompt_text,
+                    label="Enter your MarioGPT System prompt. ex: 'As an esteemed level designer renowned for creating some of the top 100 levels in Super Mario Maker...'",
+                )
+
+            text_prompt = gr.Textbox(
+                value="Generate a level with a few pipes, many coins. make sure there are only 10 enemies. Make sure there is a ground path Mario can walk on",
+                label="Enter your MarioGPT prompt. ex: 'Generate a level with a few pipes, many coins. make sure there are only 10 enemies. Make sure there is a ground path Mario can walk on'",
+            )
+
+            model = gr.Radio(
+                [
+                    # "Cladue Opus", # no opus for demo
+                    "Claude Sonnet",
+                    "Claude Haiku",
+                    "Llama3 70B",
+                    "Llama3 8B",
+                ],
+                label="Select Model",
+                value="Claude Sonnet",
+            )
+
+    # with gr.Accordion(label="Advanced settings", open=False):
+    #     temperature = gr.Number(
+    #         value=0.7,
+    #         label="temperature: Increase for more randomness",
+    #     )
+    # level_size = gr.Slider(
+    #     value=1,
+    #     minimum=1,
+    #     maximum=5,
+    #     step=1,
+    #     label="level_size",
+    # )
+
+    btn = gr.Button("Generate level")
+    with gr.Row():
+        with gr.Group():
+            level_play = gr.HTML()
+        level_image = gr.Image()
+    btn.click(
+        fn=generate,
+        inputs=[
+            # temperature,
+            # level_size,
+            model,
+            text_prompt,
+            system_prompt,
+        ],
+        outputs=[level_image, level_play],
+    )
+    gr.Examples(
+        examples=[
+            [
+                "Claude Sonnet",
+                "Generate a level with a few pipes, many coins. make sure there are only 10 enemies. Make sure there is a ground path Mario can walk on",
+            ],
+            [
+                "Claude Sonnet",
+                "Design a level with blocks arranged in a pyramid-like shape, with coins scattered around the base and goombas guarding the top. Have a pipe at the top.",
+            ],
+            [
+                "Claude Sonnet",
+                "Make a simple level that has no enemies, but lots and lots of coins. Lots of blocks for mario to walk on.",
+            ],
+        ],
+        inputs=[model, text_prompt, system_prompt_text],
+        outputs=[level_image, level_play],
+        fn=generate,
+        cache_examples=True,
+    )
+
+app.mount("/static", StaticFiles(directory="static", html=True), name="static")
+app = gr.mount_gradio_app(app, demo, "/")
+uvicorn.run(app, host="0.0.0.0", port=7860)

mario_gpt/__init__.py

@@ -0,0 +1,17 @@
+from mario_gpt.dataset import MarioDataset
+from mario_gpt.lm import MarioBert, MarioGPT, MarioLM
+from mario_gpt.prompter import Prompter
+from mario_gpt.sampler import GPTSampler, SampleOutput
+from mario_gpt.trainer import MarioGPTTrainer, TrainingConfig
+
+__all__ = [
+    "Prompter",
+    "MarioDataset",
+    "MarioBert",
+    "MarioGPT",
+    "MarioLM",
+    "SampleOutput",
+    "GPTSampler",
+    "TrainingConfig",
+    "MarioGPTTrainer",
+]

mario_gpt/data/tiles/mma_tiles.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f6d58bb3228bcd3c653c4a58b69044588ffd6e5e4c946a860497a39d84eb60b8
+size 6586

mario_gpt/dataset.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+from typing import List, Optional
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
+
+from mario_gpt.level import FULL_LEVEL_STR_WITH_PATHS
+
+DEFAULT_MODEL = "distilgpt2"
+
+
+def split_given_size(a, size):
+    return np.split(a, np.arange(size, len(a), size))
+
+
+def flip_and_transpose(arr: np.array, flip_first: bool = False):
+    if arr.shape[-1] > 1:
+        if flip_first:
+            return np.flip(arr, -1).transpose()
+        return np.flip(arr.transpose(), -1)
+    return arr
+
+
+def join_list_of_list(str_lists):
+    return ["".join(s) for s in str_lists]
+
+
+def characterize(str_lists):
+    return [list(s) for s in str_lists]
+
+
+class MarioDataset(Dataset):
+    def __init__(
+        self,
+        tokenizer: Optional[PreTrainedTokenizer] = None,
+        level_string: Optional[str] = None,
+        context_len: int = 700,
+        height: int = 14,
+        remove_start_end_tokens: bool = False,
+        sample_all_indices: bool = False,
+    ):
+        if level_string is None:
+            print(
+                "No level string specified, using default string FULL_LEVEL_STR_WITH_PATHS..."
+            )
+            level_string = FULL_LEVEL_STR_WITH_PATHS
+        elif ".txt" in level_string:
+            with open(level_string, "r") as file:
+                level_string = file.read()
+
+        self.character_set = set(level_string)
+        if "\n" in self.character_set:
+            self.character_set.remove("\n")
+        self.vocab_size = len(self.character_set)
+        self.sample_all_indices = sample_all_indices
+
+        def get_training_corpus():
+            yield list(level_string)
+
+        if tokenizer is None:
+            tokenizer = AutoTokenizer.from_pretrained(DEFAULT_MODEL)
+
+        self.tokenizer = tokenizer
+        if getattr(tokenizer, "train_new_from_iterator", None) is not None:
+            self.tokenizer = self.tokenizer.train_new_from_iterator(
+                get_training_corpus(), 52000
+            )
+        elif getattr(tokenizer, "train_from_iterator", None) is not None:
+            self.tokenizer = PreTrainedTokenizerFast(tokenizer_object=self.tokenizer)
+            self.tokenizer = self.tokenizer.train_new_from_iterator(
+                get_training_corpus(), self.vocab_size
+            )
+        self.context_len = context_len
+        self.height = height
+
+        x, self.str_arr = self.convert_level_to_tensor(level_string.split("\n"))
+        self.input_ids = x["input_ids"].squeeze()
+        self.attention_masks = x["attention_mask"].squeeze()
+        if remove_start_end_tokens:
+            self.input_ids = self.input_ids[1:-1]
+            self.attention_masks = self.attention_masks[1:-1]
+
+        self.indices = self.generate_indices()
+
+        self.unique_tokens, self.unique_counts = self.input_ids.unique(
+            return_counts=True
+        )
+        self.weighted_unique_counts = (
+            1.0 / self.unique_counts / torch.sum(self.unique_counts)
+        )
+
+        self.token_dict = {}
+        string_tokens = list(self.tokenizer.decode(self.unique_tokens))
+        for int_token, string_token in zip(self.unique_tokens, string_tokens):
+            self.token_dict[string_token] = int_token
+
+    def convert_level_to_tensor(self, level: List[str]):
+        str_arr = flip_and_transpose(np.array(characterize(level)))
+        str_arr = "".join(join_list_of_list(str_arr))
+
+        x = self.tokenizer(str_arr, return_tensors="pt")
+        return x, str_arr
+
+    def __len__(self):
+        return self.indices.shape[0]
+
+    def __getitem__(self, idx):
+        if isinstance(idx, int):
+            indices = self.indices[idx]
+        else:
+            indices = torch.stack([self.indices[i] for i in idx])
+        return self.input_ids[indices], self.attention_masks[indices]
+
+    def generate_indices(self):
+        out = []
+        for idx in range(self.input_ids.shape[0] - self.context_len):
+            if idx % self.height == 0 or self.sample_all_indices:
+                arange = torch.arange(idx, idx + self.context_len)
+                out.append(arange)
+        return torch.stack(out)
+
+    def sample_indices(self, batch_size):
+        out = []
+        for _ in range(batch_size):
+            start_idx = np.random.randint(0, self.__len__() - self.context_len)
+            indices = torch.arange(start_idx, start_idx + self.context_len)
+            out.append(indices)
+        return torch.stack(out)
+
+    def __str__(self):
+        str_list = characterize(self.tokenizer.batch_decode(self.x["input_ids"]))
+        string = "\n".join(
+            join_list_of_list(flip_and_transpose(np.array(str_list), True))
+        )
+        return string

mario_gpt/lm/__init__.py

@@ -0,0 +1,44 @@
+from typing import Optional, Union
+
+from transformers import PreTrainedModel, PreTrainedTokenizer
+
+# lm stuff
+from mario_gpt.lm.base import BaseMarioLM
+from mario_gpt.lm.bert import MarioBert
+from mario_gpt.lm.gpt import MarioGPT
+from mario_gpt.prompter import Prompter
+
+
+def MarioLM(
+    lm: Optional[PreTrainedModel] = None,
+    tokenizer: Optional[PreTrainedTokenizer] = None,
+    context_len: int = 700,
+    prompter: Optional[Prompter] = None,
+    mask_proportion: float = 0.15,
+    mask_model: bool = False,
+    lm_path: Optional[str] = None,
+    tokenizer_path: Optional[str] = None,
+    **kwargs
+) -> Union[MarioGPT, MarioBert]:
+    if not mask_model:
+        return MarioGPT(
+            lm=lm,
+            tokenizer=tokenizer,
+            context_len=context_len,
+            prompter=prompter,
+            lm_path=lm_path,
+            tokenizer_path=tokenizer_path,
+            **kwargs
+        )
+    return MarioBert(
+        lm=lm,
+        tokenizer=tokenizer,
+        context_len=context_len,
+        mask_proportion=mask_proportion,
+        lm_path=lm_path,
+        tokenizer_path=tokenizer_path,
+        **kwargs
+    )
+
+
+__all__ = ["BaseMarioLM", "MarioGPT", "MarioBert", "MarioLM"]

mario_gpt/lm/base.py

@@ -0,0 +1,91 @@
+import abc
+import os
+from typing import Any, Dict, Optional
+
+import torch
+from transformers import PreTrainedModel, PreTrainedTokenizer
+
+
+class BaseMarioLM(metaclass=abc.ABCMeta):
+
+    PRETRAINED_LM_PATH = ""
+    PRETRAINED_TOKENIZER_PATH = ""
+
+    BASE_LM_PATH = ""
+    BASE_TOKENIZER_PATH = ""
+
+    def __init__(
+        self,
+        lm: Optional[PreTrainedModel] = None,
+        tokenizer: Optional[PreTrainedTokenizer] = None,
+        context_len: int = 700,
+        lm_path: Optional[str] = None,
+        tokenizer_path: Optional[str] = None,
+        lm_kwargs: Dict[str, Any] = {},
+        tokenizer_kwargs: Dict[str, Any] = {},
+    ):
+        self.load_pretrained(
+            lm_path, tokenizer_path, lm, tokenizer, lm_kwargs, tokenizer_kwargs
+        )
+        self.context_len = context_len
+
+    def train(self):
+        self.lm.train()
+
+    def eval(self):
+        self.lm.eval()
+
+    @property
+    def device(self):
+        return self.lm.device
+
+    def to(self, device: torch.device):
+        self.lm = self.lm.to(device)
+        return self
+
+    def save_model(self, checkpoint_path: str, it: int):
+        self.lm.save_pretrained(os.path.join(checkpoint_path, f"iteration_{it}"))
+
+    @abc.abstractmethod
+    def load_pretrained_lm(
+        self, path: str, lm_kwargs: Dict[str, Any]
+    ) -> PreTrainedModel:
+        """
+        Model to be used in level tile prediction
+        """
+
+    @abc.abstractmethod
+    def load_pretrained_tokenizer(
+        self, path: str, tokenizer_kwargs: Dict[str, Any]
+    ) -> PreTrainedTokenizer:
+        """
+        Tokenizer to be used to read / decode levels
+        """
+
+    def load_pretrained(
+        self,
+        lm_path: Optional[str] = None,
+        tokenizer_path: Optional[str] = None,
+        lm: Optional[PreTrainedModel] = None,
+        tokenizer: Optional[PreTrainedTokenizer] = None,
+        lm_kwargs: Dict[str, Any] = {},
+        tokenizer_kwargs: Dict[str, Any] = {},
+    ):
+        self.lm = lm
+        self.tokenizer = tokenizer
+
+        if lm is None:
+            if lm_path is None:
+                lm_path = self.PRETRAINED_LM_PATH
+
+            print(f"Using {lm_path} lm")
+            self.lm = self.load_pretrained_lm(lm_path, lm_kwargs)
+
+        if tokenizer is None:
+            if tokenizer_path is None:
+                tokenizer_path = self.PRETRAINED_LM_PATH
+
+            print(f"Using {tokenizer_path} tokenizer")
+            self.tokenizer = self.load_pretrained_tokenizer(
+                tokenizer_path, tokenizer_kwargs
+            )

mario_gpt/lm/bert.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+import numpy as np
+import torch
+from transformers import (
+    AutoModelForMaskedLM,
+    AutoTokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    RobertaModel,
+    RobertaTokenizer,
+)
+
+from mario_gpt.lm.base import BaseMarioLM
+
+PRETRAINED_MODEL_PATH = "shyamsn97/MarioBert-448-inpaint-context-length"
+
+
+class MarioBert(BaseMarioLM):
+    PRETRAINED_LM_PATH = PRETRAINED_MODEL_PATH
+    PRETRAINED_TOKENIZER_PATH = PRETRAINED_MODEL_PATH
+
+    BASE_LM_PATH = "distilroberta-base"
+    BASE_TOKENIZER_PATH = "distilroberta-base"
+
+    def __init__(
+        self,
+        lm: Optional[PreTrainedModel] = None,
+        tokenizer: Optional[PreTrainedTokenizer] = None,
+        context_len: int = 448,
+        mask_proportion: float = 0.16,
+        lm_path: Optional[str] = None,
+        tokenizer_path: Optional[str] = None,
+        lm_kwargs: Dict[str, Any] = {},
+        tokenizer_kwargs: Dict[str, Any] = {},
+    ):
+        super().__init__(
+            lm,
+            tokenizer,
+            context_len,
+            lm_path,
+            tokenizer_path,
+            lm_kwargs,
+            tokenizer_kwargs,
+        )
+        self.mask_proportion = mask_proportion
+        self.mask_portion = int(self.context_len * self.mask_proportion)
+
+    def sample_mask(self, input_ids):
+        batch_size = input_ids.shape[0]
+        seq_len = input_ids.shape[-1]
+        mask_portion = self.mask_portion
+        sampled_start_idx = [i for i in range(seq_len - mask_portion) if i % 14 == 0]
+        sampled_start_idx = np.random.choice(sampled_start_idx, batch_size)
+        sampled_masks = []
+        for idx in sampled_start_idx:
+            mask = torch.arange(idx, idx + mask_portion)
+            sampled_masks.append(mask)
+        sampled_mask_indices = torch.stack(sampled_masks)
+        return self.apply_mask(input_ids, sampled_mask_indices)
+
+    def generate_mask(self, mask_len: int, batch_size: int = 1):
+        mask_token = self.tokenizer("<mask>").input_ids[1]
+        ones = torch.ones((batch_size, mask_len))
+        return ones * mask_token
+
+    def apply_mask(self, level, masked_indices, mask=None):
+        if len(level.shape) == 1:
+            level = level.unsqueeze(0)
+        batch_size = level.shape[0]
+        mask_len = masked_indices.shape[-1]
+        if mask is None:
+            mask = self.generate_mask(mask_len, batch_size)
+        mask = mask.long().to(level.device)
+        masked_level = level * torch.ones_like(level).to(level.device)
+        masked_level[:, masked_indices] = mask
+        return masked_level
+
+    def generate_seed(self, length: int, batch_size: Optional[int] = None):
+        seed = self.tokenizer("X", return_tensors="pt").input_ids.squeeze()[
+            1:-1
+        ]  # remove start and end tokens
+        if batch_size is None:
+            return seed.repeat(length)
+        return seed.view(1, 1).repeat(batch_size, length)
+
+    def load_pretrained_lm(self, path: str, lm_kwargs: Dict[str, Any]) -> RobertaModel:
+        return AutoModelForMaskedLM.from_pretrained(path, **lm_kwargs)
+
+    def load_pretrained_tokenizer(
+        self, path: str, tokenizer_kwargs: Dict[str, Any]
+    ) -> RobertaTokenizer:
+        return AutoTokenizer.from_pretrained(path, **tokenizer_kwargs)

mario_gpt/lm/gpt.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+import torch
+from transformers import (
+    AutoConfig,
+    AutoModelWithLMHead,
+    AutoTokenizer,
+    GPT2Model,
+    GPT2Tokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+
+from mario_gpt.lm.base import BaseMarioLM
+from mario_gpt.prompter import Prompter
+from mario_gpt.sampler import GPTSampler, SampleOutput
+
+PRETRAINED_MODEL_PATH = "shyamsn97/Mario-GPT2-700-context-length"
+
+
+class MarioGPT(BaseMarioLM):
+    PRETRAINED_LM_PATH = PRETRAINED_MODEL_PATH
+    PRETRAINED_TOKENIZER_PATH = PRETRAINED_MODEL_PATH
+
+    BASE_LM_PATH = "distilgpt2"
+    BASE_TOKENIZER_PATH = "distilgpt2"
+
+    def __init__(
+        self,
+        lm: Optional[PreTrainedModel] = None,
+        tokenizer: Optional[PreTrainedTokenizer] = None,
+        context_len: int = 700,
+        prompter: Optional[Prompter] = None,
+        lm_path: Optional[str] = None,
+        tokenizer_path: Optional[str] = None,
+        lm_kwargs: Dict[str, Any] = {},
+        tokenizer_kwargs: Dict[str, Any] = {},
+    ):
+        super().__init__(
+            lm,
+            tokenizer,
+            context_len,
+            lm_path,
+            tokenizer_path,
+            lm_kwargs,
+            tokenizer_kwargs,
+        )
+        self.prompter = prompter
+        if prompter is None:
+            self.prompter = Prompter(self.tokenizer)
+
+    def generate_seed(self, length: int, batch_size: Optional[int] = None):
+        seed = self.tokenizer("X", return_tensors="pt").input_ids.squeeze()
+        if batch_size is None:
+            return seed.repeat(length)
+        return seed.view(1, 1).repeat(batch_size, length)
+
+    def load_pretrained_lm(self, path: str, lm_kwargs: Dict[str, Any]) -> GPT2Model:
+        if path == "random":
+            print("Initializing random weights...")
+            config = AutoConfig.from_pretrained(
+                self.BASE_LM_PATH, **{**lm_kwargs, "add_cross_attention": True}
+            )
+            return AutoModelWithLMHead.from_config(config)
+        return AutoModelWithLMHead.from_pretrained(
+            path, **{**lm_kwargs, "add_cross_attention": True}
+        )
+
+    def load_pretrained_tokenizer(
+        self, path: str, tokenizer_kwargs: Dict[str, Any]
+    ) -> GPT2Tokenizer:
+        if path == "random":
+            return AutoTokenizer.from_pretrained(
+                self.BASE_TOKENIZER_PATH, **tokenizer_kwargs
+            )
+        return AutoTokenizer.from_pretrained(path, **tokenizer_kwargs)
+
+    def sample(
+        self,
+        seed: Optional[torch.Tensor] = None,
+        prompts: Optional[List[str]] = None,
+        num_steps: int = 1,
+        temperature: float = 2.0,
+        encoder_hidden_states: torch.Tensor = None,
+        use_tqdm: bool = False,
+        return_tensor: bool = False,
+    ) -> SampleOutput:
+        sampler = GPTSampler(self, temperature, 16, self.context_len, use_tqdm)
+        return sampler(
+            seed=seed,
+            prompts=prompts,
+            num_steps=num_steps,
+            encoder_hidden_states=encoder_hidden_states,
+            return_tensor=return_tensor,
+        )

mario_gpt/prompter.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+import random
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from scipy import stats
+from transformers import pipeline
+
+from mario_gpt.dataset import MarioDataset
+from mario_gpt.utils import view_level
+
+STATISTICS = {
+    "enemy": np.array([1.0, 3.0, 7.0]),
+    "pipe": np.array([0.0, 2.0, 5.0]),
+    "block": np.array([50.0, 75.0, 176.0]),
+}
+
+FEATURE_EXTRACTION_MODEL = "facebook/bart-base"
+
+
+class Prompter:
+    def __init__(
+        self,
+        level_tokenizer,
+        prompter_model: str = FEATURE_EXTRACTION_MODEL,
+        use_raw_counts: bool = False,
+        statistics: Optional[Dict[str, Any]] = None,
+    ):
+        self.prompter_model = prompter_model
+        self.feature_extraction = pipeline(
+            "feature-extraction",
+            model=prompter_model,
+            tokenizer=prompter_model,
+            framework="pt",
+        )
+
+        self.level_tokenizer = level_tokenizer
+
+        self.use_raw_counts = use_raw_counts
+        self.statistics = statistics
+        if statistics is None:
+            self.statistics = STATISTICS
+
+    @property
+    def pipe_thresholds(self) -> Tuple[List[int], List[str]]:
+        thresholds = self.statistics["pipe"]
+        keywords = ["no", "little", "some", "many"]
+        return thresholds, keywords
+
+    @property
+    def enemy_thresholds(self) -> Tuple[List[int], List[str]]:
+        thresholds = self.statistics["enemy"]
+        keywords = ["no", "little", "some", "many"]
+        return thresholds, keywords
+
+    @property
+    def block_thresholds(self) -> Tuple[List[int], List[str]]:
+        thresholds = self.statistics["block"]
+        keywords = ["little", "little", "some", "many"]
+        return thresholds, keywords
+
+    def count_pipes(self, flattened_level: str) -> int:
+        return flattened_level.count("<>")
+
+    def count_enemies(self, flattened_level: str) -> int:
+        return flattened_level.count("E") + flattened_level.count("B")
+
+    def count_blocks(self, flattened_level: str) -> int:
+        return np.sum([flattened_level.count(char) for char in ["X", "S", "?", "Q"]])
+
+    def _flatten_level(self, string_level: List[str]) -> str:
+        return "".join(string_level)
+
+    def pipe_prompt(self, flattened_level: str, level: str) -> str:
+        count = self.count_pipes(flattened_level)
+        keyword = f"{count}"
+        if not self.use_raw_counts:
+            thresholds, keywords = self.pipe_thresholds
+            threshold = np.digitize(count, thresholds, right=True)
+            keyword = keywords[threshold]
+        return f"{keyword} pipes", keyword
+
+    def enemy_prompt(self, flattened_level: str, level: str) -> str:
+        count = self.count_enemies(flattened_level)
+        keyword = f"{count}"
+        if not self.use_raw_counts:
+            thresholds, keywords = self.enemy_thresholds
+            threshold = np.digitize(count, thresholds, right=True)
+            keyword = keywords[threshold]
+        return f"{keyword} enemies", keyword
+
+    def block_prompt(self, flattened_level: str, level: str) -> str:
+        count = self.count_blocks(flattened_level)
+        keyword = f"{count}"
+        if not self.use_raw_counts:
+            thresholds, keywords = self.block_thresholds
+            threshold = np.digitize(count, thresholds, right=True)
+            keyword = keywords[threshold]
+        return f"{keyword} blocks", keyword
+
+    def elevation_prompt(self, flattened_level: str, level: str):
+        top_levels = level[:6]  # elevation 8 and up
+        for t in top_levels:
+            if "X" in t or "<" in t or ">" in t:
+                return "high elevation", "high"
+        return "low elevation", "low"
+
+    def output_hidden(self, prompt: str, device: torch.device = torch.device("cpu")):
+        # Reducing along the first dimension to get a 768 dimensional array
+        return (
+            self.feature_extraction(prompt, return_tensors="pt")[0]
+            .mean(0)
+            .to(device)
+            .view(1, -1)
+        )
+
+    def dataset_statistics(self, dataset: MarioDataset):
+        enemy_counts = []
+        pipe_counts = []
+        block_counts = []
+        for i in range(len(dataset)):
+            level, _ = dataset[i]
+            str_level = self._flatten_level(view_level(level, dataset.tokenizer))
+
+            enemy_count = self.count_enemies(str_level)
+            pipe_count = self.count_pipes(str_level)
+            block_count = self.count_blocks(str_level)
+
+            enemy_counts.append(enemy_count)
+            pipe_counts.append(pipe_count)
+            block_counts.append(block_count)
+        d = {"enemy": {}, "pipe": {}, "block": {}}
+
+        d["enemy"] = stats.mstats.mquantiles(enemy_counts, [0.33, 0.66, 0.95])
+        d["pipe"] = stats.mstats.mquantiles(pipe_counts, [0.33, 0.66, 0.95])
+        d["block"] = stats.mstats.mquantiles(block_counts, [0.33, 0.66, 0.95])
+        return d
+
+    def __call__(
+        self, level: torch.Tensor = None, sample_prompt: bool = False
+    ) -> Union[str, torch.Tensor]:
+        device: torch.device = torch.device("cpu")
+        if not sample_prompt:
+            if level is None:
+                raise ValueError("Level must be provided if sample_prompt is not true!")
+            str_level = view_level(level, self.level_tokenizer)
+            flattened_level = self._flatten_level(str_level)
+
+            pipe_prompt, _ = self.pipe_prompt(flattened_level, str_level)
+            enemy_prompt, _ = self.enemy_prompt(flattened_level, str_level)
+            block_prompt, _ = self.block_prompt(flattened_level, str_level)
+            elevation_prompt, _ = self.elevation_prompt(flattened_level, str_level)
+            device = level.device
+        else:
+            str_level = None
+            pipe_prompt = random.choice(["no", "little", "some", "many"]) + " pipes"
+            enemy_prompt = random.choice(["no", "little", "some", "many"]) + " enemies"
+            block_prompt = (
+                random.choice(["little", "little", "some", "many"]) + " blocks"
+            )  # levels always have blocks
+            elevation_prompt = (
+                random.choice(["low", "high"]) + " elevation"
+            )  # levels always have blocks
+
+        prompt_dict = {
+            "pipe": pipe_prompt,
+            "enemy": enemy_prompt,
+            "block": block_prompt,
+            "elevation_prompt": elevation_prompt,
+        }
+        prompt = f"{pipe_prompt}, {enemy_prompt}, {block_prompt}, {elevation_prompt}"
+        hidden = self.output_hidden(prompt, device=device)
+        return prompt, hidden, prompt_dict, str_level

mario_gpt/sampler.py

@@ -0,0 +1,370 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from PIL.Image import Image
+from tqdm import tqdm
+from transformers import LogitsProcessorList, TemperatureLogitsWarper, TopKLogitsWarper
+
+from mario_gpt.lm.base import BaseMarioLM
+from mario_gpt.prompter import Prompter
+from mario_gpt.simulator import Simulator
+from mario_gpt.utils import (
+    convert_level_to_png,
+    load_level,
+    save_level,
+    trim_level,
+    view_level,
+)
+
+
+@dataclass
+class SampleOutput:
+    level: Optional[List[str]]
+    prompt: Optional[str] = None
+    img: Optional[Image] = None
+    sample_predictions_str: Optional[List[str]] = None
+    sample_predictions_img: Optional[Image] = None
+    level_tensor: Optional[torch.Tensor] = None
+    sample_predictions_tensor: Optional[torch.Tensor] = None
+
+    @classmethod
+    def create(
+        cls,
+        level_tensor: torch.Tensor,
+        sample_predictions_tensor: torch.Tensor,
+        tokenizer,
+        prompter: Optional[Prompter] = None,
+    ) -> SampleOutput:
+        # batch = 1
+        level = None
+        img = None
+
+        try:
+            level = view_level(level_tensor, tokenizer)
+            img = convert_level_to_png(level)[0]
+        except Exception as e:
+            print(
+                f"Failed to generate string or image representation for full level! Got error {e}"
+            )
+            level = None
+            img = None
+        try:
+            sample_predictions_str = view_level(sample_predictions_tensor, tokenizer)
+            sample_predictions_img = convert_level_to_png(sample_predictions_str)[0]
+        except Exception as e:
+            print(
+                f"Failed to generate string or image representation for sampled predictions! Got error {e}"
+            )
+            sample_predictions_str = None
+            sample_predictions_img = None
+
+        prompt = None
+        if prompter is not None:
+            prompt = prompter(level_tensor)[0]
+
+        return SampleOutput(
+            level,
+            prompt,
+            img,
+            sample_predictions_str,
+            sample_predictions_img,
+            level_tensor,
+            sample_predictions_tensor,
+        )
+
+    @classmethod
+    def from_level_predictions(
+        cls,
+        level: torch.Tensor,
+        sample_predictions: torch.Tensor,
+        tokenizer,
+        prompter: Optional[Prompter] = None,
+    ) -> Union[SampleOutput, List[SampleOutput]]:
+        level_tensor = trim_level(level).squeeze().detach().cpu()
+        sample_predictions_tensor = (
+            trim_level(sample_predictions).squeeze().detach().cpu()
+        )
+
+        if len(level_tensor.shape) == 1:
+            return SampleOutput.create(
+                level_tensor, sample_predictions_tensor, tokenizer, prompter
+            )
+
+        out = []
+        for _level_tensor, _sample_predictions_tensor in zip(
+            level_tensor, sample_predictions_tensor
+        ):
+            sample_output = SampleOutput.create(
+                _level_tensor, _sample_predictions_tensor, tokenizer, prompter
+            )
+            out.append(sample_output)
+        return out
+
+    def save(self, filename: str) -> str:
+        save_level(self.level, filename)
+
+    @classmethod
+    def load(cls, filename: str) -> SampleOutput:
+        level = load_level(filename)
+        return SampleOutput(level=level)
+
+    def play(self):
+        simulator = Simulator(level=self.level)
+        simulator.interactive()
+
+    def run_astar(self, render=True):
+        simulator = Simulator(level=self.level)
+        simulator.astar(render)
+
+
+class GPTSampler:
+    def __init__(
+        self,
+        mario_lm: BaseMarioLM,
+        temperature: float = 2.0,
+        top_k: int = 16,
+        context_len: int = 700,
+        use_tqdm: bool = False,
+        use_argmax: bool = False,
+    ):
+        self.mario_lm = mario_lm
+        self.temperature = temperature
+        self.top_k = top_k
+        self.context_len = context_len
+        self.use_tqdm = use_tqdm
+        self.use_argmax = use_argmax
+        self.logits_processor = LogitsProcessorList()
+        self.logits_warper = LogitsProcessorList(
+            [
+                TopKLogitsWarper(top_k),  # number of characters
+                TemperatureLogitsWarper(temperature),
+            ]
+        )
+
+    @property
+    def device(self) -> torch.device:
+        return self.mario_lm.device
+
+    def step(
+        self,
+        seed: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        with torch.no_grad():
+            attention_mask = torch.ones_like(seed).to(seed.device)
+            input_ids = seed
+            out = self.mario_lm.lm(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                encoder_hidden_states=encoder_hidden_states,
+                token_type_ids=None,
+            )
+            logits = out.logits.detach()
+            if len(logits.shape) == 2:
+                logits = logits.view(1, 1, -1)
+            next_token_logits = logits[:, -1, :]
+
+            if self.use_argmax:
+                next_tokens = next_token_logits.argmax(-1)
+            else:
+                next_token_scores = self.logits_processor(input_ids, next_token_logits)
+                next_token_scores = self.logits_warper(input_ids, next_token_scores)
+                probs = torch.nn.functional.softmax(next_token_scores, dim=-1)
+                next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+        return next_tokens, encoder_hidden_states
+
+    def sample(
+        self,
+        seed: Union[Optional[torch.Tensor], Optional[SampleOutput]] = None,
+        prompts: Optional[List[str]] = None,
+        num_steps: int = 1,
+        encoder_hidden_states: torch.Tensor = None,
+        return_tensor: bool = False,
+    ):
+        self.mario_lm.eval()
+        context_len = self.context_len - 28
+        with torch.no_grad():
+            if seed is None:
+                seed = self.mario_lm.generate_seed(1, batch_size=len(prompts)).to(
+                    self.device
+                )
+                out_tensor = seed.to(self.device)
+            elif isinstance(seed, SampleOutput):
+                out_tensor = seed.level_tensor.to(self.device).squeeze()
+            else:
+                out_tensor = seed.to(self.device).squeeze()
+            if len(out_tensor.shape) < 2:
+                # if we pass in a single seed vector, then we repeat for each prompt
+                # Otherwise, we treat inputs as separate seed-prompt pairs
+                out_tensor = out_tensor.view(1, -1).repeat(len(prompts), 1)
+            if encoder_hidden_states is None:
+                if prompts is not None:
+                    encoder_hidden_states = torch.stack(
+                        [
+                            self.mario_lm.prompter.output_hidden(prompt)
+                            for prompt in prompts
+                        ]
+                    )
+                else:
+                    encoder_hidden_states = torch.stack(
+                        [
+                            self.mario_lm.prompter(sample_prompt=True)[1]
+                            for _ in range(seed.shape[0])
+                        ]
+                    )
+            encoder_hidden_states = encoder_hidden_states.to(
+                self.device
+            )  # b x 1 x hidden_dim
+            encoder_hidden_states = encoder_hidden_states.view(
+                out_tensor.shape[0], 1, -1
+            )
+            if not self.use_tqdm:
+                bar = np.arange(num_steps)
+            else:
+                bar = tqdm(np.arange(num_steps))
+            with torch.no_grad():
+                for i in bar:
+                    inp = out_tensor * 1
+                    if len(out_tensor.shape) > 0 and out_tensor.shape[-1] > context_len:
+                        diff = inp.shape[-1] % 14  # height of mario level
+                        ctx = context_len + diff
+                        inp = inp[:, -ctx:] * 1
+                    next_tokens, encoder_hidden_states = self.step(
+                        inp,
+                        encoder_hidden_states=encoder_hidden_states,
+                    )
+                    out_tensor = torch.cat(
+                        [out_tensor, next_tokens.unsqueeze(-1)], dim=-1
+                    )
+                    if self.use_tqdm:
+                        bar.set_description(
+                            f"shape: {inp.shape}, {out_tensor.shape} first: {inp[0][0]}, last: {out_tensor[0][-1]}"
+                        )
+            if self.use_tqdm:
+                bar.close()
+        sample_out = SampleOutput.from_level_predictions(
+            out_tensor,
+            out_tensor[:, -num_steps:],
+            self.mario_lm.tokenizer,
+            self.mario_lm.prompter,
+        )
+        self.mario_lm.train()
+        if return_tensor:
+            return sample_out, out_tensor
+        return sample_out
+
+    def __call__(self, *args, **kwargs):
+        return self.sample(*args, **kwargs)
+
+
+class BertSampler:
+    def __init__(
+        self,
+        mario_lm: BaseMarioLM,
+        temperature: float = 2.0,
+        top_k: int = 16,
+        context_len: int = 448,
+        mask_proportion: float = 0.16,
+    ):
+        self.mario_lm = mario_lm
+        self.temperature = temperature
+        self.top_k = top_k
+        self.logits_processor = LogitsProcessorList()
+        self.logits_warper = LogitsProcessorList(
+            [
+                TopKLogitsWarper(top_k),  # number of characters
+                TemperatureLogitsWarper(temperature),
+            ]
+        )
+        self.context_len = context_len
+        self.mask_proportion = mask_proportion
+        self.mask_portion = int(self.context_len * self.mask_proportion)
+        self.mask_portion = self.mask_portion - self.mask_portion % 14 + 14
+
+    @property
+    def device(self) -> torch.device:
+        return self.mario_lm.device
+
+    def get_context(self, input_ids, mask_indices):
+        start_idx = mask_indices[0]
+        end_idx = mask_indices[-1]
+
+        if input_ids.shape[-1] <= self.context_len:
+            clipped = input_ids.shape[-1] % 14
+            input_ids = input_ids[:clipped]
+
+        portion = (self.context_len - self.mask_portion) / 2
+
+        remainder = 0
+        left = start_idx - portion
+        if left < 0:
+            remainder = -1 * left
+
+        right = end_idx + portion + remainder
+
+        return input_ids[left:right]
+
+    def sample(
+        self,
+        seed: Union[torch.Tensor, SampleOutput],
+        mask: torch.Tensor,
+        return_tensor: bool = False,
+    ):
+        self.mario_lm.eval()
+        mask_indices = mask.nonzero()
+        input_ids = seed
+        if isinstance(seed, SampleOutput):
+            input_ids = seed.level_tensor.to(self.device).squeeze()
+
+        input_id_list = []
+        for i in range(input_ids.shape[0]):
+            input_id = input_ids[i]
+            mask_index = mask_indices[mask_indices[:, 0] == i][:, -1]
+            input_id = self.get_context(input_id, mask_index)
+            input_id_list.append(input_id)
+        input_ids = torch.stack(input_ids, dim=0).to(self.device)
+
+        attention_mask = torch.ones_like(input_ids).to(seed.device)
+
+        if len(input_ids.shape) < 2:
+            # if we pass in a single seed vector, then we repeat for each prompt
+            # Otherwise, we treat inputs as separate seed-prompt pairs
+            input_ids = input_ids.view(1, -1)
+
+        out = self.mario_lm.lm(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=None,
+        )
+        logits = out.logits.detach()
+        if len(logits.shape) == 2:
+            logits = logits.view(1, 1, -1)
+
+        if self.use_argmax:
+            tokens = logits.argmax(-1)
+        else:
+            tokens_scores = self.logits_processor(input_ids, tokens)
+            tokens_scores = self.logits_warper(input_ids, tokens_scores)
+            probs = torch.nn.functional.softmax(tokens_scores, dim=-1)
+            tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+
+        out = input_ids.detach()
+
+        for i in range(input_ids.shape[0]):
+            mask_index = mask_indices[mask_indices[:, 0] == i][:, -1]
+            out[i, mask_index] = tokens[i, mask_index].detach()
+
+        sample_out = SampleOutput.from_level_predictions(
+            out,
+            tokens,
+            self.mario_lm.tokenizer,
+            self.mario_lm.prompter,
+        )
+        self.mario_lm.train()
+        if return_tensor:
+            return sample_out, tokens
+        return sample_out

mario_gpt/simulator/__init__.py

@@ -0,0 +1,3 @@
+from mario_gpt.simulator.simulator import Simulator
+
+__all__ = ["Simulator"]