plancraft 0.1.4__py3-none-any.whl → 0.3.0__py3-none-any.whl

plancraft/environments/sampler.py

@@ -1,224 +0,0 @@
-import math
-import random
-from collections import Counter
-
-import numpy as np
-from plancraft.environments.items import all_data, ALL_ITEMS
-from plancraft.environments.recipes import RECIPES
-from plancraft.environments.planner import optimal_planner, get_ancestors
-
-
-MAX_STACK_SIZE = {}
-for data_item in all_data["items"]:
-    if data_item["stackable"]:
-        MAX_STACK_SIZE[data_item["type"]] = data_item["stackSize"]
-    else:
-        MAX_STACK_SIZE[data_item["type"]] = 1
-
-
-def sample_distractors(
-    exclude_set: set = None, num_distractors: int = 16
-) -> dict[str, int]:
-    distractors = {}
-    while len(distractors) < num_distractors:
-        item = random.choice(ALL_ITEMS)
-        if exclude_set is not None and item in exclude_set:
-            continue
-        count = random.randint(1, MAX_STACK_SIZE[item])
-        distractors[item] = count
-    return distractors
-
-
-def assign_to_slots(inventory: dict[str, int]) -> list[dict]:
-    # slots available outside of crafting interface
-    available_slots = list(range(10, 46))
-    random.shuffle(available_slots)
-    inventory_list = []
-
-    for item, total_count in inventory.items():
-        while total_count > 0:
-            if len(available_slots) == 0:
-                print("Not enough slots available")
-                break
-            slot = available_slots.pop()
-            count_in_slot = min(total_count, MAX_STACK_SIZE[item])
-            inventory_list.append({"slot": slot, "item": item, "count": count_in_slot})
-            total_count -= count_in_slot
-
-    return inventory_list
-
-
-def sample_recipes(
-    target: str,
-    overall_exclude_set: set,
-    target_count: int = 1,
-    current_depth=0,
-    max_depth=20,
-) -> tuple[set, set]:
-    # stop if the depth is too high
-    if current_depth > max_depth:
-        return {}, overall_exclude_set
-
-    # get all the recipes that can craft the target
-    overall_exclude_set.update([target])
-    local_exclude_set = set()
-    random_recipes = []
-    for r in RECIPES[target]:
-        recipe_inputs, exclude_set = r.sample_inputs()
-        # if inputs are already in the exclude set, skip this recipe (ensures no cycle)
-        if exclude_set.intersection(overall_exclude_set):
-            return {}, overall_exclude_set
-        local_exclude_set.update(exclude_set)
-        random_recipes.append((r, recipe_inputs))
-
-    overall_exclude_set |= local_exclude_set
-
-    # no recipes found
-    if len(random_recipes) == 0:
-        return {}, overall_exclude_set
-
-    # sample a random recipe
-    random_recipe = random.choice(random_recipes)
-    recipe, start_inputs = random_recipe
-
-    # recipe will not produce enough
-    if recipe.result.count < target_count:
-        # must do recipe X times
-        recipe_multiplier = math.ceil(target_count / recipe.result.count)
-        start_inputs = {k: v * recipe_multiplier for k, v in start_inputs.items()}
-
-    for input_item in list(start_inputs.keys()):
-        # randomize depth first search to end early
-        if random.choice([True, False]):
-            continue
-
-        children_recipe_inputs, updated_exclude_set = sample_recipes(
-            target=input_item,
-            overall_exclude_set=overall_exclude_set,
-            target_count=start_inputs[input_item],
-            current_depth=current_depth + 1,
-        )
-        if len(children_recipe_inputs) == 0:
-            continue
-
-        overall_exclude_set.update(updated_exclude_set)
-
-        # remove recipe input item since we are crafting it
-        start_inputs[input_item] = 0
-
-        # add the children recipe inputs
-        for item, count in children_recipe_inputs.items():
-            start_inputs[item] = start_inputs.get(item, 0) + count
-
-    overall_exclude_set = overall_exclude_set - {None}
-    start_inputs = {k: v for k, v in start_inputs.items() if v > 0}
-
-    return start_inputs, overall_exclude_set
-
-
-def remove_ancestor_items(target: str, inventory: dict[str, int]) -> dict[str, int]:
-    ancestors = set(get_ancestors(target))
-    possible_items = set(inventory.keys())
-    items_to_remove = list(ancestors.intersection(possible_items))
-    num_items = random.randint(1, len(items_to_remove))
-    for item in random.sample(items_to_remove, num_items):
-        count_to_remove = random.randint(1, inventory[item])
-        inventory[item] -= count_to_remove
-        if inventory[item] == 0:
-            del inventory[item]
-    return inventory
-
-
-def construct_example(
-    target: str,
-    num_distractors: 16,
-    impossible=False,
-) -> list[dict]:
-    """
-    For a given target object, number of distractors, and impossible flag
-    Return a dictionary with the start inventory for the crafting task
-
-    The crafting task should be to craft the target, the inventory should contain
-    the resources required for the recipe to be crafted.
-
-    The number of distractors are how many random items should be added to the inventory.
-
-    If impossible is True, the target item should not be craftable with the given inventory.
-    """
-
-    # sample the recipe
-    inventory, overall_exclude_set = sample_recipes(target, set())
-    if impossible:
-        # if impossible then remove one or more items from the inventory
-        inventory = remove_ancestor_items(
-            target,
-            inventory,
-        )
-
-    # add distractors to the inventory
-    distractors = sample_distractors(overall_exclude_set, num_distractors)
-    inventory.update(distractors)
-
-    optimal_path = optimal_planner(target, inventory)
-    # @TODO this is a hack to ensure that we don't have impossible examples
-    while optimal_path is not None and impossible:
-        inventory = remove_ancestor_items(target, inventory)
-        optimal_path = optimal_planner(target, inventory)
-
-    # assign to slots
-    inventory_list = assign_to_slots(inventory)
-    example = {
-        "inventory": inventory,
-        "slotted_inventory": inventory_list,
-        "target": target,
-        "num_distractors": num_distractors,
-        "impossible": impossible,
-    }
-    # either impossible and no path or not impossible and path exists
-    assert (impossible and optimal_path is None) or (
-        not impossible and optimal_path is not None
-    )
-
-    if not impossible:
-        example["optimal_path_length"] = len(optimal_path)
-        example["optimal_path"] = [r.result.item for (r, i) in optimal_path]
-        example["inventory_trace"] = [i for (r, i) in optimal_path]
-        items_used, unique_items_used = calculate_stats_from_inventory_trace(
-            [example["inventory"]] + example["inventory_trace"]
-        )
-        example["items_used"] = items_used
-        example["unique_items_used"] = unique_items_used
-
-    return example
-
-
-def calculate_stats_from_inventory_trace(
-    inventory_trace: list[dict],
-) -> tuple[int, int]:
-    total_items_used = 0
-    total_unique_items_used = 0
-
-    for a, b in zip(inventory_trace[:-1], inventory_trace[1:]):
-        diff = Counter(a) - Counter(b)
-        total_items_used += sum(diff.values())
-        total_unique_items_used += len(diff)
-
-    return total_items_used, total_unique_items_used
-
-
-def generate_dataset(seed=2024, distractors=[4, 8, 16], num_examples=10):
-    random.seed(seed)
-    np.random.seed(seed)
-
-    dataset = []
-    for recipe_target in list(RECIPES.keys()):
-        if len(RECIPES[recipe_target]) == 0:
-            continue
-        for num_distractors in distractors:
-            for _ in range(num_examples):
-                example = construct_example(
-                    target=recipe_target, num_distractors=num_distractors
-                )
-                dataset.append(example)
-
-    return dataset

plancraft/models/few_shot_images/__init__.py

@@ -1,16 +0,0 @@
-import os
-import glob
-
-import numpy as np
-import imageio
-
-
-def get_few_shot_images_path():
-    return os.path.dirname(__file__)
-
-
-def load_prompt_images() -> list[np.ndarray]:
-    current_dir = get_few_shot_images_path()
-    files = glob.glob(os.path.join(current_dir, "*.png"))
-    images = [imageio.imread(file) for file in files]
-    return images

plancraft/models/oam.py

@@ -1,283 +0,0 @@
-from typing import Optional
-
-from loguru import logger
-
-import torch
-import torch.nn as nn
-import torchvision.transforms.v2 as v2
-from transformers import (
-    AutoConfig,
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    PretrainedConfig,
-    PreTrainedModel,
-)
-
-from plancraft.models.bbox_model import IntegratedBoundingBoxModel
-
-
-class PlancraftOAMConfig(PretrainedConfig):
-    model_type = "plancraft-aom"
-    is_composition = True
-
-    def __init__(
-        self,
-        from_llama=False,
-        **kwargs,
-    ):
-        self.from_llama = from_llama
-        super().__init__(**kwargs)
-
-
-class PlancraftOAM(PreTrainedModel):
-    config_class = PlancraftOAMConfig
-
-    def __init__(self, config: PlancraftOAMConfig):
-        super().__init__(config)
-
-        self.config = config
-        # load text model
-        if self.config.from_llama:
-            self.text_model = AutoModelForCausalLM.from_pretrained(
-                "meta-llama/Meta-Llama-3.1-8B-Instruct",
-            )
-        else:
-            text_model_config = AutoConfig.from_pretrained(
-                "meta-llama/Meta-Llama-3.1-8B-Instruct",
-            )
-            self.text_model = AutoModelForCausalLM.from_config(text_model_config)
-
-        # load vision model
-        self.vision_model = IntegratedBoundingBoxModel.from_pretrained(
-            "gautierdag/plancraft-maskrcnn"
-        )
-        self.vision_model.eval()
-
-        # convert vision features to text embedding
-        self.vision_to_text_embedding = nn.Linear(
-            1024, self.text_model.config.hidden_size
-        )
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            "meta-llama/Meta-Llama-3.1-8B-Instruct",
-            trust_remote=True,
-        )
-        # add special tokens
-        self.tokenizer.add_special_tokens(
-            {
-                "additional_special_tokens": [
-                    "<|inventory|>",
-                ]
-            }
-        )
-        self.tokenizer.pad_token = self.tokenizer.eos_token
-        self.inventory_idx = self.tokenizer.convert_tokens_to_ids("<|inventory|>")
-
-        # resize token embeddings
-        self.text_model.resize_token_embeddings(len(self.tokenizer))
-        # image transforms
-        self.transforms = v2.Compose(
-            [v2.ToImage(), v2.ToDtype(torch.float32, scale=True)]
-        )
-
-    @torch.no_grad()
-    def extract_bboxes(self, images: list) -> list[dict]:
-        if len(images) == 0:
-            return []
-        img_tensors = torch.stack([self.transforms(img) for img in images])
-        img_tensors = img_tensors.cuda()
-        # disable gradients
-        self.vision_model.freeze()
-        # get bounding box predictions
-        bbox_preds = self.vision_model(img_tensors)
-        return bbox_preds
-
-    def prepare_messages(self, messages: list, bboxes: list[dict]) -> str:
-        # no bounding boxes
-        if len(bboxes) == 0:
-            text = self.tokenizer.apply_chat_template(
-                messages, add_generation_prompt=not self.training, tokenize=False
-            )
-            text = text.replace("<|begin_of_text|>", "")
-            return text
-
-        # expand <|inventory|> tokens into N tokens (N = number of bounding boxes)
-        new_messages = []
-        i_pred = 0
-        for m in messages:
-            new_message = m.copy()
-            if new_message["role"] == "user" and new_message["content"].endswith(
-                "<|inventory|>"
-            ):
-                # add inventory tokens for each bounding box
-                new_message["content"] = new_message["content"].replace(
-                    "<|inventory|>",
-                    "<|inventory|>" * (bboxes[i_pred]["features"].shape[0]),
-                )
-                i_pred += 1
-            new_messages.append(new_message)
-
-        assert i_pred == len(
-            bboxes
-        ), "Number of inventory tokens does not match number of images"
-        # add special tokens
-
-        text = self.tokenizer.apply_chat_template(
-            new_messages, add_generation_prompt=not self.training, tokenize=False
-        )
-        text = text.replace("<|begin_of_text|>", "")
-        return text
-
-    def inputs_merger(
-        self,
-        input_ids: torch.LongTensor,
-        inputs_embeds: Optional[torch.Tensor],
-        image_hidden_states: Optional[torch.Tensor],
-    ):
-        # along batch dimension
-        for i in range(len(image_hidden_states)):
-            if len(image_hidden_states[i]) == 0:
-                assert (
-                    input_ids[i] == self.inventory_idx
-                ).sum() == 0, "No images but inventory token is still present"
-                continue
-
-            # count the number of inventory tokens
-            n_inventory_tokens = (input_ids[i] == self.inventory_idx).sum()
-            if n_inventory_tokens != image_hidden_states[i].shape[0]:
-                logger.warning(
-                    f"Number of inventory tokens ({n_inventory_tokens}) does not match number of bounding boxes ({image_hidden_states[i].shape[0]}). Possible truncation."
-                )
-                # truncated from the start
-                image_hidden_states[i] = image_hidden_states[i][-n_inventory_tokens:]
-
-            # replace inventory tokens with bbox features
-            inputs_embeds[i, input_ids[i] == self.inventory_idx] = image_hidden_states[
-                i
-            ]
-        return inputs_embeds
-
-    def process_inputs(
-        self,
-        batch_messages: list[list[dict]] = [],  # list of list of messages (untokenized)
-        batch_images: list[list] = [],  # list of list of images (unprocessed)
-    ) -> tuple[dict[str, torch.FloatTensor], list[torch.FloatTensor], int]:
-        """
-        Converts raw images and messages into model inputs
-        """
-        assert len(batch_images) == len(
-            batch_messages
-        ), "Number of images and messages should match in the batch dim"
-        # initial forward pass
-        texts_batch = []
-        image_hidden_states = []
-        total_boxes = 0
-        for images, messages in zip(batch_images, batch_messages):
-            # process images
-            bboxes = self.extract_bboxes(images)
-            if len(bboxes) > 0:
-                # get bbox features
-                features = torch.concat([p["features"] for p in bboxes], dim=0)
-                # upscale to text embedding size
-                features_embeds = self.vision_to_text_embedding(features)
-                image_hidden_states.append(features_embeds)
-                # count bboxes total
-                total_boxes += features.shape[0]
-            else:
-                image_hidden_states.append([])
-
-            # process messages
-            text = self.prepare_messages(messages, bboxes)
-            texts_batch.append(text)
-
-        # tokenize text
-        # @NOTE: truncation could cause issues with inventory tokens not matching number of boxes
-        # in that case, we will truncate the boxes from the end, and issue a warning
-        batch = self.tokenizer(
-            texts_batch,
-            truncation=True,
-            padding=True,
-            max_length=16384,
-            return_tensors="pt",
-        )
-        return batch, image_hidden_states, total_boxes
-
-    def forward(
-        self,
-        batch_messages: list[list[dict]] = [],  # list of list of messages (untokenized)
-        batch_images: list[list] = [],  # list of list of images (unprocessed)
-        **kwargs,
-    ):
-        labels = None
-        batch, image_hidden_states, total_boxes = self.process_inputs(
-            batch_messages, batch_images
-        )
-        # move to cuda
-        batch = {k: v.cuda() for k, v in batch.items()}
-        attention_mask = batch["attention_mask"]
-        input_ids = batch["input_ids"]
-
-        labels = input_ids.clone()
-        # remove inventory tokens from labels
-        labels[labels == self.inventory_idx] = -100
-        # sanity check: should have same number of boxes as inventory tokens
-        assert (labels == -100).sum() == total_boxes
-
-        # get text embeddings
-        inputs_embeds = self.text_model.get_input_embeddings()(input_ids)
-        inputs_embeds = self.inputs_merger(
-            input_ids, inputs_embeds, image_hidden_states
-        )
-        # forward pass
-        return self.text_model(
-            inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            labels=labels,
-            return_dict=True,
-        )
-
-    @torch.no_grad()
-    def generate(
-        self,
-        batch_messages: list[list[dict]],
-        batch_images: list[list],
-        do_sample=True,
-        temperature=0.6,
-        max_new_tokens=32,
-    ):
-        self.tokenizer.padding_side = "left"
-
-        batch, image_hidden_states, _ = self.process_inputs(
-            batch_messages, batch_images
-        )
-        batch = {k: v.cuda() for k, v in batch.items()}
-        attention_mask = batch["attention_mask"]
-        input_ids = batch["input_ids"]
-
-        inputs_embeds = self.text_model.get_input_embeddings()(input_ids)
-        inputs_embeds = self.inputs_merger(
-            input_ids, inputs_embeds, image_hidden_states
-        )
-
-        generated_sequences = self.text_model.generate(
-            inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            do_sample=do_sample,
-            temperature=temperature,
-            max_new_tokens=max_new_tokens,
-            pad_token_id=self.tokenizer.eos_token_id,
-        )
-
-        # Decode the output
-        text_responses = self.tokenizer.batch_decode(
-            generated_sequences,
-            # generated_sequences[:, prompt_tokens:],
-            skip_special_tokens=False,
-        )
-
-        # remove <|eot_id|> tokens
-        text_responses = [
-            text_response.replace("<|eot_id|>", "") for text_response in text_responses
-        ]
-
-        _, total_tokens_used = generated_sequences.shape
-        return text_responses, total_tokens_used

plancraft/models/prompts.py

@@ -1,158 +0,0 @@
-from plancraft.models.utils import gold_search_recipe
-
-VALID_ACTIONS = ["move", "smelt", "think", "search", "impossible"]
-
-ACTIONS_DESCRIPTIONS = {
-    "move": {
-        "description": "Transfer a specific quantity of an item from one slot to another",
-        "format": "`move: from [Source] to [Target] with quantity N`",
-    },
-    "smelt": {
-        "description": "Smelt an item in a furnace and moves the output to a specific slot",
-        "format": "`smelt: from [Source] to [Target] with quantity N`",
-    },
-    "think": {
-        "description": "Generate thoughts to help you decide on the next action",
-        "format": "`think: <thought message>`",
-    },
-    "search": {
-        "description": "Search for a recipe to craft a specific item",
-        "format": "`search: <recipe name>`",
-    },
-    "impossible": {
-        "description": "Stop task if it is certain that it is impossible with given inventory",
-        "format": "`impossible: <reason>`",
-    },
-}
-
-BASE_SYSTEM_PROMPT = """You are crafting in Minecraft. You need to decide on the next action.
-
-Crafting Grid: The crafting table is organized into a 3x3 grid. Each slot in the grid has a unique identifier:
-    - Top row: [A1] [A2] [A3]
-    - Middle row: [B1] [B2] [B3]
-    - Bottom row: [C1] [C2] [C3]
-
-The output of the crafting process is placed in a designated output slot labeled [0] You cannot move or smelt items directly into slot [0]
-
-Inventory Slots: The remaining inventory slots (outside of the crafting grid) are used for storing items. These slots are labeled as [I1] to [I36]"""
-
-BASE_SYSTEM_PROMPT_EXAMPLE = """Example:
-    - `move: from [I2] to [A1] with quantity 3`
-    - `smelt: from [I5] to [I6] with quantity 1`
-
-Constraints:
-   - You cannot move or smelt items into [0]
-   - If an item is not in slot [0] then the recipe is incorrect
-   - You need to move items from [0] to a free inventory slot to complete the crafting process"""
-
-
-def get_system_prompt(actions: list[str]):
-    assert set(actions).issubset(VALID_ACTIONS), f"Invalid actions: {actions}"
-    assert "move" in actions, "move should be one of the actions"
-    assert "smelt" in actions, "smelt should be one of the actions"
-
-    descriptions = ""
-    for action in actions:
-        descriptions += f"\n\t- {action}: {ACTIONS_DESCRIPTIONS[action]['description']}"
-
-    output_format = ""
-    for action in actions:
-        output_format += f"\n\t- {ACTIONS_DESCRIPTIONS[action]['format']}"
-
-    return f"{BASE_SYSTEM_PROMPT}\n\nActions:{descriptions}\n\nFormat{output_format}\n\n{BASE_SYSTEM_PROMPT_EXAMPLE}"
-
-
-CRAFTING_STEPS = [
-    "Craft an item of type: andesite\ninventory:\n - diorite [I18] quantity 1\n - cobblestone [I30] quantity 1",
-    "Craft an item of type: andesite\ninventory:\n - diorite [B1] quantity 1\n - cobblestone [I30] quantity 1",
-    "Craft an item of type: andesite\ninventory:\n - andesite [0] quantity 1\n - diorite [B1] quantity 1\n - cobblestone [B2] quantity 1",
-    "Craft an item of type: iron_ingot\ninventory:\n - iron_ore [I36] quantity 1\n - cobblestone [I30] quantity 1",
-]
-
-BASE_ACTION_STEPS = [
-    "move: from [I18] to [B1] with quantity 1",
-    "move: from [I30] to [B2] with quantity 1",
-    "move: from [0] to [I6] with quantity 1",
-    "smelt: from [I36] to [I35] with quantity 1",
-]
-
-THINK_STEPS = [
-    "think: To solve this task I need to craft andesite using 1 diorite and 1 cobblestone side by side.",
-    "think: Now I need to move the cobblestone into position [B2] to be right of the diorite.",
-    "think: Now I can craft the andesite by moving it from the craft slot [0] to a free inventory slot.",
-    "think: To craft an iron_ingot, I need to smelt iron_ore into an empty slot.",
-]
-
-SEARCH_STEPS = [
-    "search: andesite",
-    None,
-    None,
-    "search: iron_ingot",
-]
-
-
-def get_prompt_example(
-    actions: list[str],
-    use_text_inventory=True,
-    use_multimodal_content_format=False,
-    use_images=False,
-) -> list[dict]:
-    assert set(actions).issubset(VALID_ACTIONS), f"Invalid actions: {actions}"
-    assert "move" in actions, "move should be one of the actions"
-    assert "smelt" in actions, "smelt should be one of the actions"
-
-    if use_images:
-        assert (
-            use_multimodal_content_format
-        ), "use_images requires use_multimodal_content_format"
-
-    example_dialogue = []
-    for i, step in enumerate(CRAFTING_STEPS):
-        text = step
-        if not use_text_inventory:
-            text = text.split("\ninventory:\n")[0]
-
-        example_dialogue.append({"role": "user", "content": text})
-        if "search" in actions and SEARCH_STEPS[i]:
-            example_dialogue.append({"role": "assistant", "content": SEARCH_STEPS[i]})
-            search_target = text.split("seach: ")[-1].strip()
-            search_response = gold_search_recipe(search_target)
-            example_dialogue.append({"role": "user", "content": search_response})
-        if "think" in actions:
-            example_dialogue.append({"role": "assistant", "content": THINK_STEPS[i]})
-            example_dialogue.append({"role": "user", "content": "Ok"})
-        example_dialogue.append({"role": "assistant", "content": BASE_ACTION_STEPS[i]})
-
-    if not use_multimodal_content_format:
-        return example_dialogue
-
-    # convert to multimodal dialogue
-    multimodal_dialogue = []
-    for message in example_dialogue:
-        if "Craft an item" in message["content"]:
-            content_list = [
-                {
-                    "type": "text",
-                    "text": message["content"],
-                }
-            ]
-            if use_images:
-                content_list.append(
-                    {
-                        "type": "image",
-                    }
-                )
-
-            multimodal_dialogue.append(
-                {"role": message["role"], "content": content_list}
-            )
-        else:
-            multimodal_dialogue.append(
-                {
-                    "role": message["role"],
-                    "content": [
-                        {"type": "text", "text": message["content"]},
-                    ],
-                }
-            )
-    return multimodal_dialogue