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

plancraft/models/act.py

@@ -1,67 +1,46 @@
-import copy
 import torch
-from dotenv import load_dotenv
 
 from plancraft.config import EvalConfig
-from plancraft.environments.actions import (
-    NoOp,
-    StopAction,
-    SymbolicAction,
-)
-from plancraft.models.base import ABCModel, History
+from plancraft.models.base import PlancraftBaseModel
 from plancraft.models.bbox_model import IntegratedBoundingBoxModel
-from plancraft.models.few_shot_images import load_prompt_images
 from plancraft.models.generators import (
-    OAMGenerator,
     OpenAIGenerator,
     TransformersGenerator,
 )
-from plancraft.models.prompts import get_prompt_example, get_system_prompt
-from plancraft.models.utils import (
-    convert_observation_to_message,
-    parse_content_response,
-)
 
+from plancraft.utils import History
 
-load_dotenv()
 
-
-class ActModel(ABCModel):
+class ActModel(PlancraftBaseModel):
     """
     Model that does action without thinking step
     """
 
     def __init__(self, cfg: EvalConfig):
         self.cfg = cfg
-        self.env_is_multimodal = not cfg.plancraft.environment.symbolic
-        self.use_maskrcnn = cfg.plancraft.use_maskrcnn
+        self.use_fasterrcnn = cfg.plancraft.use_fasterrcnn
         self.use_multimodal_content_format = cfg.plancraft.use_multimodal_content_format
         self.use_text_inventory = cfg.plancraft.use_text_inventory
         self.use_images = cfg.plancraft.use_images
 
         self.bbox_model = None
-        if self.use_maskrcnn:
-            assert self.env_is_multimodal, "MaskRCNN only supported in multimodal mode"
+        if self.use_fasterrcnn:
+            # fasterrcnn is not multimodal model but a separate model
             self.bbox_model = IntegratedBoundingBoxModel.from_pretrained(
-                "gautierdag/plancraft-maskrcnn"
+                "gautierdag/plancraft-fasterrcnn"
             )
             self.bbox_model.eval()
             if torch.cuda.is_available():
                 self.bbox_model.cuda()
-            # MaskRCNN is not multimodal model but a separate model
-
-        self.few_shot = cfg.plancraft.few_shot
-        self.use_system_prompt = cfg.plancraft.system_prompt
-        self.max_invalid_actions = 3
 
         # underlying language model
         if "gpt-4o" in cfg.plancraft.model:
             self.use_multimodal_content_format = True
             self.llm = OpenAIGenerator(
-                use_images=self.use_images, model_name=cfg.plancraft.model
+                use_images=self.use_images,
+                model_name=cfg.plancraft.model,
+                api_key=cfg.env_variables.openai_api_key,
             )
-        elif "oam" in cfg.plancraft.model:
-            self.llm = OAMGenerator(model_name=cfg.plancraft.model)
         else:
             # model is transformers based
             self.llm = TransformersGenerator(
@@ -70,116 +49,26 @@ class ActModel(ABCModel):
                 quantize=cfg.plancraft.quantize,
                 use_hot_cache=cfg.plancraft.hot_cache,
                 adapter_name=cfg.plancraft.adapter,
+                hf_token=cfg.env_variables.hf_token,
             )
-
-        self.prompt_images = []
-
-        self.valid_actions = cfg.plancraft.valid_actions
-        self.system_prompt_text = get_system_prompt(self.valid_actions)
-
-        examples = get_prompt_example(
-            self.valid_actions,
-            use_text_inventory=self.use_text_inventory,
-            use_multimodal_content_format=self.use_multimodal_content_format,
-            use_images=self.use_images,
-        )
-        if self.env_is_multimodal and self.use_images:
-            self.prompt_images = load_prompt_images()
-
-        if self.use_multimodal_content_format:
-            self.system_prompt = {
-                "role": "system",
-                "content": [
-                    {"text": copy.deepcopy(self.system_prompt_text), "type": "text"}
-                ],
-            }
-        else:
-            self.system_prompt = {
-                "role": "system",
-                "content": copy.deepcopy(self.system_prompt_text),
-            }
-
-        if not self.few_shot:
-            examples = []
-        if not self.use_system_prompt:
-            self.system_prompt = None
-
-        self._history = History(
-            initial_dialogue=examples,
-            use_multimodal_content_format=self.use_multimodal_content_format,
-        )
-
         self.max_messages_window = cfg.plancraft.max_message_window
         self.kv_cache = None
 
-    @property
-    def history(self):
-        return self._history
-
-    def reset_history(
-        self,
-        objective: str,
-    ):
-        examples = []
-        if self.few_shot:
-            examples = get_prompt_example(
-                self.valid_actions,
-                use_text_inventory=self.use_text_inventory,
-                use_multimodal_content_format=self.use_multimodal_content_format,
-                use_images=self.use_images,
-            )
-
-        self.history.reset(objective=objective, initial_dialogue=examples)
+    def reset(self):
         self.llm.reset()
 
-    def step(self, observation: dict) -> SymbolicAction | StopAction:
-        self.history.add_observation_to_history(observation)
-
-        # add observation to history
-        observation_message = convert_observation_to_message(
-            observation,
-            objective=self.history.objective,
-            bbox_model=self.bbox_model,
-            oam_model="oam" in self.llm.model_name,
-            use_text_inventory=self.use_text_inventory,
-            use_multimodal_content_format=self.use_multimodal_content_format,
-            use_images=self.use_images,
+    def step(self, observation: dict, dialogue_history: History) -> str:
+        # get message window
+        message_window, image_window = self.llm.prepare_messages(
+            history=dialogue_history,
+            max_messages_window=self.max_messages_window,
         )
-        self.history.add_message_to_history(content=observation_message, role="user")
-
-        # Iterate until valid action
-        i = 0
-        while i < self.max_invalid_actions:
-            # add observation to history
-            message_window, image_window = self.llm.prepare_messages(
-                history=self.history,
-                max_messages_window=self.max_messages_window,
-                system_prompt=self.system_prompt,
-                prompt_images=self.prompt_images,
-            )
-            action_messages, action_token_used = self.llm.generate_unconstrained(
-                batch_messages=[message_window],
-                images=[image_window],
-            )
-            self.history.tokens_used += action_token_used
-
-            action_message = action_messages[0].split("\n")[0].strip()
-
-            self.history.add_message_to_history(
-                content=action_message, role="assistant"
-            )
-            response = parse_content_response(
-                action_message, valid_actions=self.valid_actions
-            )
-            if not isinstance(response, str):
-                # valid action
-                self.history.add_action_to_history(response)
-                return response
-
-            self.history.add_message_to_history(
-                content=response,
-            )
-            i += 1
-
-        # if no action is found after max_invalid_actions, default to useless move action
-        return NoOp()
+        # generate next action
+        action_messages, action_token_used = self.llm.generate_unconstrained(
+            batch_messages=[message_window],
+            images=[image_window],
+        )
+        # update tokens used
+        dialogue_history.tokens_used += action_token_used
+        # return raw action message
+        return action_messages[0].split("\n")[0].strip()

plancraft/models/base.py

@@ -1,161 +1,24 @@
 import abc
-from collections import Counter
-from copy import copy
 
-from loguru import logger
+from plancraft.utils import History
 
-from plancraft.environments.actions import (
-    SymbolicMoveAction,
-    SymbolicSmeltAction,
-)
 
-
-class History:
-    def __init__(
-        self,
-        objective: str = "",
-        initial_dialogue: list[dict] = [],
-        use_multimodal_content_format=False,
-    ):
-        self.dialogue_history = initial_dialogue
-        self.initial_dialogue_length = len(initial_dialogue)
-        self.action_history = []
-        self.inventory_history = []
-        self.inventory_counters = []
-        self.images = []
-        self.objective = objective
-        self.tokens_used = 0
-        self.use_multimodal_content_format = use_multimodal_content_format
-
-    def add_message_to_history(self, content: str | dict, role="user"):
-        if role == "assistant":
-            logger.info(content)
-
-        if isinstance(content, dict):
-            assert "content" in content, "content key not found in message"
-            content["role"] = role
-            self.dialogue_history.append(content)
-        else:
-            # fix for listed content type
-            if self.use_multimodal_content_format:
-                return self.add_message_to_history(
-                    content={
-                        "content": [{"type": "text", "text": content}],
-                        "role": role,
-                    },
-                    role=role,
-                )
-            else:
-                self.dialogue_history.append({"role": role, "content": content})
-
-    def add_action_to_history(self, action: SymbolicSmeltAction | SymbolicMoveAction):
-        if action is None:
-            return
-        self.action_history.append(action.model_dump())
-
-    def add_inventory_to_history(self, inventory: list[dict[str, int]]):
-        self.inventory_history.append(inventory)
-
-        # count inventory
-        counter = Counter()
-        for item in inventory:
-            # ignore slot 0
-            if "slot" in item and item["slot"] == 0:
-                continue
-            if "index" in item and item["index"] == 0:
-                continue
-            counter[item["type"]] += item["quantity"]
-
-        self.inventory_counters.append(counter)
-
-    def add_image_to_history(self, image):
-        self.images.append(image)
-
-    def add_observation_to_history(self, observation: dict):
-        if observation is None:
-            return
-        if "inventory" in observation:
-            clean_inv = []
-            # remove empty slots
-            for item in observation["inventory"]:
-                if item["quantity"] > 0:
-                    clean_inv.append(item)
-            self.add_inventory_to_history(clean_inv)
-        if "pov" in observation:
-            self.add_image_to_history(observation["pov"])
-
-    def __str__(self):
-        return str(self.dialogue_history)
-
-    def reset(self, objective: str = "", initial_dialogue: list[dict] = []):
-        self.dialogue_history = initial_dialogue
-        self.action_history = []
-        self.inventory_history = []
-        self.inventory_counters = []
-        self.images = []
-        self.objective = objective
-
-    def set_objective(self, objective: str):
-        self.objective = objective
-
-    def trace(self):
-        return {
-            "dialogue_history": copy(
-                self.dialogue_history[self.initial_dialogue_length :]
-            ),
-            "action_history": copy(self.action_history),
-            "inventory_history": copy(self.inventory_history),
-            "objective": copy(self.objective),
-            "tokens_used": copy(self.tokens_used),
-        }
-
-    @property
-    def num_steps(self):
-        return len(self.action_history)
-
-    def check_stuck(self, max_steps_no_change: int = 10) -> bool:
-        """
-        If inventory content does not change for max_steps_no_change steps
-        the agent is considered stuck.
-
-        With N=10, the oracle solver can still solve 100% of the examples
-        """
-        if len(self.inventory_counters) <= max_steps_no_change:
-            return False
-
-        return all(
-            c == self.inventory_counters[-max_steps_no_change - 1]
-            for c in self.inventory_counters[-max_steps_no_change - 1 :]
-        )
-
-
-class ABCModel(abc.ABC):
+class PlancraftBaseModel(abc.ABC):
     """
     Model class must implement the following methods to work with evaluator
     """
 
     @abc.abstractmethod
-    def step(
-        self, observation: list[dict]
-    ) -> list[SymbolicMoveAction | SymbolicSmeltAction]:
+    def step(self, observation: dict, dialogue_history: History) -> str:
         """
-        Model should output a valid action based on the 3 types available
-
-        Note this is a batch operation, so the model should return a list of actions
-        for each observation in the batch
+        Model should output an action in text based on the types available
+        We also pass history to the model to allow for chat models to track the dialogue
         """
         raise NotImplementedError()
 
-    @property
     @abc.abstractmethod
-    def history(self) -> History:
+    def reset(self):
         """
-        Return the history object.
+        Reset the model state - ready for a new episode
         """
         raise NotImplementedError()
-
-    def reset_history(self, objective: str = ""):
-        """
-        Reset the history with a new objective.
-        """
-        self.history.reset(objective=objective)

plancraft/models/bbox_model.py

@@ -19,7 +19,7 @@ from torchvision.models.detection.roi_heads import (
 )
 from torchvision.ops import boxes as box_ops
 
-from plancraft.environments.items import ALL_ITEMS
+from plancraft.environment.items import ALL_ITEMS
 
 from functools import lru_cache
 
@@ -455,7 +455,7 @@ class IntegratedBoundingBoxModel(nn.Module, PyTorchModelHubMixin):
             preds = self.model(x)
             return preds
 
-    def get_inventory(self, pil_image, resolution="high"):
+    def get_inventory(self, pil_image, resolution="high") -> dict:
         """
         Predict boxes and quantities
         """
@@ -467,21 +467,24 @@ class IntegratedBoundingBoxModel(nn.Module, PyTorchModelHubMixin):
         return self.prediction_to_inventory(predictions[0], resolution=resolution)
 
     @staticmethod
-    def prediction_to_inventory(
-        prediction, threshold=0.9, resolution="high"
-    ) -> list[dict]:
-        inventory = []
+    def prediction_to_inventory(prediction, threshold=0.9, resolution="high") -> dict:
+        inventory = {}
         seen_slots = set()
-        for i in range(len(prediction["boxes"])):
-            slot = bbox_to_slot_index_iou(prediction["boxes"][i], resolution=resolution)
-            score = prediction["scores"][i]
-            label_idx = prediction["labels"][i].item()
-            label = ALL_ITEMS[label_idx]
-            quantity = prediction["quantities"][i].item()
-            if score > threshold:
-                if slot in seen_slots:
-                    continue
-                inventory.append({"slot": slot, "type": label, "quantity": quantity})
+        for bbox, score, label, quantity in zip(
+            prediction["boxes"],
+            prediction["scores"],
+            prediction["labels"],
+            prediction["quantities"],
+        ):
+            slot = bbox_to_slot_index_iou(bbox, resolution=resolution)
+            if score < threshold:
+                break
+            if slot in seen_slots:
+                continue
+            label = ALL_ITEMS[label.item()]
+            quantity = quantity.item()
+            inventory["slot"] = {"type": label, "quantity": quantity}
+            seen_slots.add(slot)
         return inventory
 
     def freeze(self):

plancraft/models/dummy.py

@@ -1,51 +1,38 @@
 import random
 
 from plancraft.config import EvalConfig
-from plancraft.environments.actions import (
-    SymbolicMoveAction,
-    SymbolicSmeltAction,
+from plancraft.environment.actions import (
+    MoveAction,
 )
-from plancraft.models.base import ABCModel, History
+from plancraft.models.base import PlancraftBaseModel
 
 
-class DummyModel(ABCModel):
+class DummyModel(PlancraftBaseModel):
     """
     Dummy model returns actions that do random action
     """
 
     def __init__(self, cfg: EvalConfig):
-        self._history = History(objective="")
+        pass
 
-    @property
-    def history(self):
-        return self._history
+    def reset(self):
+        pass
 
     def random_select(self, observation):
-        if observation is None or "inventory" not in observation:
-            return SymbolicMoveAction(slot_from=0, slot_to=0, quantity=1)
         # randomly pick an item from the inventory
         item_indices = set()
-        for item in observation["inventory"]:
+        for slot, item in observation["inventory"].items():
             if item["quantity"] > 0:
-                item_indices.add(item["index"])
+                item_indices.add(slot)
         all_slots_to = set(range(1, 46))
         empty_slots = all_slots_to - item_indices
 
         random_slot_from = random.choice(list(item_indices))
         random_slot_to = random.choice(list(empty_slots))
 
-        return SymbolicMoveAction(
+        return MoveAction(
             slot_from=random_slot_from, slot_to=random_slot_to, quantity=1
         )
 
-    def step(self, observation: dict) -> list[SymbolicMoveAction | SymbolicSmeltAction]:
-        # add observation to history
-        self.history.add_observation_to_history(observation)
-
-        # get action
-        action = self.random_select(observation)
-
-        # add action to history
-        self.history.add_action_to_history(action)
-
-        return action
+    def step(self, observation: dict, **kwargs) -> str:
+        return str(self.random_select(observation))