plancraft 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

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

models/oracle.py

@@ -1,265 +0,0 @@
-import copy
-from collections import Counter
-
-from plancraft.config import EvalConfig
-from plancraft.environments.actions import (
-    RealActionInteraction,
-    StopAction,
-    SymbolicMoveAction,
-    SymbolicSmeltAction,
-)
-from plancraft.environments.planner import optimal_planner
-from plancraft.environments.recipes import (
-    ShapedRecipe,
-    ShapelessRecipe,
-    SmeltingRecipe,
-    id_to_item,
-)
-from plancraft.environments.sampler import MAX_STACK_SIZE
-from plancraft.models.base import ABCModel, History
-
-
-def item_set_id_to_type(item_set_ids: set[int]):
-    return set(id_to_item(i) for i in item_set_ids)
-
-
-def find_free_inventory_slot(inventory: list[dict], from_slot: int) -> int:
-    # find a free slot in the inventory for the item in from_slot
-    from_item_type, from_item_quantity = None, None
-
-    type_to_slot = {}
-    slot_to_quantity = {}
-    for item in inventory:
-        if ("slot" in item and item["slot"] == from_slot) or (
-            "index" in item and item["index"] == from_slot
-        ):
-            from_item_quantity = item["quantity"]
-            from_item_type = item["type"]
-            # break
-        item_type = item["type"]
-        quantity = item["quantity"]
-        if quantity == 0:
-            item_type = "air"
-
-        if "slot" in item:
-            item_slot = item["slot"]
-        else:
-            item_slot = item["index"]
-
-        if item_type not in type_to_slot:
-            type_to_slot[item_type] = [item_slot]
-        else:
-            type_to_slot[item_type].append(item_slot)
-
-        if item_slot not in slot_to_quantity:
-            slot_to_quantity[item_slot] = quantity
-        else:
-            slot_to_quantity[item_slot] += quantity
-
-    assert from_item_type is not None, f"Item not found in slot {from_slot}"
-
-    # if there is a free slot with the same item type
-    if from_item_type in type_to_slot:
-        for slot in type_to_slot[from_item_type]:
-            if (
-                slot != from_slot
-                and slot_to_quantity[slot] + from_item_quantity
-                <= MAX_STACK_SIZE[from_item_type]
-            ):
-                return slot
-
-    # if there is a free slot with air
-    for slot in type_to_slot["air"]:
-        if slot != from_slot and slot > 10:
-            return slot
-
-    raise ValueError("No free slot found")
-
-
-def find_item_in_inventory(target: str, inventory: list[dict]) -> int:
-    for item in inventory:
-        if item["type"] == target and item["quantity"] > 0:
-            if "slot" in item:
-                return item["slot"]
-            elif "index" in item:
-                return item["index"]
-            raise ValueError("Neither slot or index is set")
-
-
-def get_inventory_counter(inventory: list[dict]) -> Counter:
-    counter = Counter()
-    for item in inventory:
-        if "slot" in item and item["slot"] == 0:
-            continue
-        if "index" in item and item["index"] == 0:
-            continue
-        if item["type"] == "air":
-            continue
-        counter[item["type"]] += item["quantity"]
-    return counter
-
-
-def get_crafting_slot_item(inventory: list[dict]) -> dict:
-    for item in inventory:
-        if "slot" in item and item["slot"] == 0 and item["quantity"] > 0:
-            return item
-        if "index" in item and item["index"] == 0 and item["quantity"] > 0:
-            return item
-    return None
-
-
-def update_inventory(
-    inventory: list[dict], slot_from: int, slot_to: int, quantity: int
-) -> list[dict]:
-    """
-    decrements quantity of item in slot_from
-    NOTE: we don't care about incrementing the items in slot_to
-
-    """
-    new_inventory = []
-    for item in inventory:
-        if "slot" in item and item["slot"] == slot_from:
-            item["quantity"] -= quantity
-        elif "index" in item and item["index"] == slot_from:
-            item["quantity"] -= quantity
-        new_inventory.append(item)
-    return new_inventory
-
-
-class OracleModel(ABCModel):
-    """
-    Oracle model returns actions that solve the task optimally
-    """
-
-    def __init__(self, cfg: EvalConfig):
-        assert (
-            cfg.plancraft.environment.symbolic_action_space
-        ), "Only symbolic actions are supported for oracle"
-        self.history = History(objective="")
-        self.plans = []
-        self.subplans = []
-
-    def reset_history(self, objective: str = ""):
-        self.history.reset(objective=objective)
-        self.plans = []
-        self.subplans = []
-
-    def get_plan(self, observation: dict):
-        # objective="Craft an item of type: ...."
-        # this simply recovering the target item to craft
-        target = self.history.objective.split(": ")[-1]
-        inventory_counter = get_inventory_counter(observation["inventory"])
-        self.plans = optimal_planner(target=target, inventory=inventory_counter)
-
-    def get_next_action(
-        self, observation: dict
-    ) -> SymbolicMoveAction | SymbolicSmeltAction:
-        if len(self.subplans) > 0:
-            return self.subplans.pop(0)
-        if len(self.plans) == 0:
-            raise ValueError("No more steps in plan")
-
-        observed_inventory = copy.deepcopy(observation["inventory"])
-
-        # take item from crafting slot
-        if slot_item := get_crafting_slot_item(observed_inventory):
-            # move item from crafting slot to inventory
-            free_slot = find_free_inventory_slot(observed_inventory, from_slot=0)
-            return SymbolicMoveAction(
-                slot_from=0, slot_to=free_slot, quantity=slot_item["quantity"]
-            )
-
-        plan_recipe, new_inventory = self.plans.pop(0)
-        self.subplans = []
-        new_inventory_counter = Counter(new_inventory)
-        current_inventory = observed_inventory
-        current_inventory_counter = get_inventory_counter(current_inventory)
-        items_to_use_counter = current_inventory_counter - new_inventory_counter
-        new_items = new_inventory_counter - current_inventory_counter
-        assert len(new_items) == 1
-
-        if isinstance(plan_recipe, ShapelessRecipe):
-            crafting_slot = 1
-
-            # add each item to crafting slots
-            for item, quantity in items_to_use_counter.items():
-                n = 0
-                while n < quantity:
-                    from_slot = find_item_in_inventory(item, current_inventory)
-
-                    # skip if from_slot is the crafting slot
-                    if from_slot == crafting_slot:
-                        crafting_slot += 1
-                        n += 1
-                        continue
-
-                    # low_level_plan.append(("move", item, from_slot, crafting_slot, 1))
-                    action = SymbolicMoveAction(
-                        slot_from=from_slot, slot_to=crafting_slot, quantity=1
-                    )
-                    # update state of inventory
-                    current_inventory = update_inventory(
-                        current_inventory, from_slot, crafting_slot, 1
-                    )
-                    self.subplans.append(action)
-
-                    crafting_slot += 1
-                    n += 1
-
-        # if plan_recipe is a smelting recipe
-        elif isinstance(plan_recipe, SmeltingRecipe):
-            assert len(items_to_use_counter) == 1, "smelting only supports one item"
-            for item, quantity in items_to_use_counter.items():
-                from_slot = find_item_in_inventory(item, current_inventory)
-                free_slot = find_free_inventory_slot(
-                    current_inventory, from_slot=from_slot
-                )
-                action = SymbolicSmeltAction(
-                    slot_from=from_slot, slot_to=free_slot, quantity=quantity
-                )
-                self.subplans.append(action)
-
-        # if plan_recipe is a shaped recipe
-        elif isinstance(plan_recipe, ShapedRecipe):
-            for i, row in enumerate(plan_recipe.kernel):
-                for j, item_set in enumerate(row):
-                    inventory_position = (i * 3) + j + 1
-                    valid_items = item_set_id_to_type(item_set)
-                    for item in valid_items:
-                        if items_to_use_counter[item] > 0:
-                            from_slot = find_item_in_inventory(item, current_inventory)
-                            action = SymbolicMoveAction(
-                                slot_from=from_slot,
-                                slot_to=inventory_position,
-                                quantity=1,
-                            )
-                            items_to_use_counter[item] -= 1
-                            # update state of inventory
-                            current_inventory = update_inventory(
-                                current_inventory, from_slot, inventory_position, 1
-                            )
-                            self.subplans.append(action)
-                            break
-        else:
-            raise NotImplementedError(f"Recipe type {type(plan_recipe)} not supported")
-
-        return self.subplans.pop(0)
-
-    def step(
-        self, observation: dict
-    ) -> list[SymbolicMoveAction | RealActionInteraction | SymbolicSmeltAction]:
-        # add observation to history
-        self.history.add_observation_to_history(observation)
-
-        # get action
-        if len(self.plans) == 0:
-            self.get_plan(observation)
-            if self.plans is None:
-                self.plans = []
-                return StopAction()
-
-        action = self.get_next_action(observation)
-
-        # add action to history
-        self.history.add_action_to_history(action)
-        return action