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

plancraft/models/generators.py

@@ -1,8 +1,6 @@
-import os
 import time
 
 import torch
-from dotenv import load_dotenv
 from loguru import logger
 from openai import OpenAI
 from PIL import Image
@@ -15,15 +13,12 @@ from transformers import (
 )
 from transformers.cache_utils import DynamicCache
 
-from plancraft.models.base import History
-from plancraft.models.oam import PlancraftOAM
 from plancraft.models.utils import (
     get_downloaded_models,
     numpy_to_base64,
     tokenize,
 )
-
-load_dotenv()
+from plancraft.utils import History
 
 
 class TransformersGenerator:
@@ -35,10 +30,12 @@ class TransformersGenerator:
         use_images=False,
         use_hot_cache=True,
         adapter_name="",
+        hf_token=None,
         **kwargs,
     ):
         self.model_name = model_name
         self.use_hot_cache = use_hot_cache
+        self.hf_token = hf_token
 
         if tokenizer_name == "same":
             tokenizer_name = model_name
@@ -71,7 +68,7 @@ class TransformersGenerator:
         else:
             self.tokenizer = AutoTokenizer.from_pretrained(
                 tokenizer_name,
-                token=os.getenv("HF_TOKEN"),  # trust_remote_code=True
+                token=self.hf_token,  # trust_remote_code=True
                 padding_side="left",  # ensure that the padding is on the left
             )
             logger.info("Loading model")
@@ -162,10 +159,9 @@ class TransformersGenerator:
             for kvs in self.past_key_values_kwargs["past_key_values"]
         ]
 
-    @staticmethod
-    def build_model_kwargs(model_name: str, **kwargs) -> tuple[str, dict]:
+    def build_model_kwargs(self, model_name: str, **kwargs) -> tuple[str, dict]:
         model_kwargs = {
-            "token": os.getenv("HF_TOKEN"),
+            "token": self.hf_token,
             # "attn_implementation": "flash_attention_2",
             # "trust_remote_code": True,
         }
@@ -203,8 +199,6 @@ class TransformersGenerator:
         self,
         history: History,
         max_messages_window: int,
-        system_prompt: dict = None,
-        prompt_images: list = [],
     ) -> tuple[list[dict], list]:
         """
         Prepare the messages using a history
@@ -214,20 +208,19 @@ class TransformersGenerator:
         if len(message_window) > 0 and message_window[0]["role"] == "assistant":
             message_window = message_window[1:]
         # add the system prompt if the first message is not a system message
-        if message_window[0]["role"] != "system" and system_prompt is not None:
-            message_window = [system_prompt] + message_window
+        if message_window[0]["role"] != "system":
+            message_window = [history.system_prompt_dialogue] + message_window
 
         image_window = []
         if self.use_images:
-            image_list = prompt_images + history.images
             image_count = 0
             # iterate through the messages in reverse order to assign images
             for m in message_window:
                 for content in m["content"]:
                     if content["type"] == "image":
                         image_count += 1
-            assert image_count <= len(image_list), "Too many images"
-            image_window = image_list[-image_count:]
+            assert image_count <= len(history.images), "Too many images"
+            image_window = history.images[-image_count:]
             image_window = [Image.fromarray(img) for img in image_window]
 
         return message_window, image_window
@@ -309,8 +302,8 @@ class TransformersGenerator:
 
 
 class OpenAIGenerator:
-    def __init__(self, use_images=False, model_name="gpt-4o-mini"):
-        self.client = OpenAI()
+    def __init__(self, use_images=False, model_name="gpt-4o-mini", api_key=None):
+        self.client = OpenAI(api_key=api_key)
         self.use_images = use_images
         self.model_name = model_name
 
@@ -321,8 +314,6 @@ class OpenAIGenerator:
         self,
         history: History,
         max_messages_window: int,
-        system_prompt: dict = None,
-        prompt_images: list = [],
     ) -> tuple[list[dict], list]:
         """
         Prepare the image messages for the model
@@ -332,11 +323,10 @@ class OpenAIGenerator:
         if len(message_window) > 0 and message_window[0]["role"] == "assistant":
             message_window = message_window[1:]
         # add the system prompt if the first message is not a system message
-        if message_window[0]["role"] != "system" and system_prompt is not None:
-            message_window = [system_prompt] + message_window
+        if message_window[0]["role"] != "system":
+            message_window = [history.system_prompt_dialogue] + message_window
 
         if self.use_images:
-            image_list = prompt_images + history.images
             img_idx = -1
             seen_images = 0
             # iterate through the messages in reverse order to assign images
@@ -346,7 +336,7 @@ class OpenAIGenerator:
                     if content["type"] == "text":
                         new_content_list.append(content)
                     elif content["type"] == "image":
-                        base64_image = numpy_to_base64(image_list[img_idx])
+                        base64_image = numpy_to_base64(history.images[img_idx])
                         img_idx -= 1
                         seen_images + 1
                         new_content = {
@@ -357,7 +347,7 @@ class OpenAIGenerator:
                         }
                         new_content_list.append(new_content)
                     message_window[i]["content"] = new_content_list
-            assert seen_images <= len(image_list), "Too many images"
+            assert seen_images <= len(history.images), "Too many images"
 
         return message_window, []
 
@@ -365,6 +355,7 @@ class OpenAIGenerator:
         self,
         batch_messages: list[list[dict]],
         max_tokens=256,
+        temperature=0.6,
         **kwargs,
     ) -> tuple[list[str], int]:
         contents = []
@@ -373,7 +364,7 @@ class OpenAIGenerator:
             response = self.client.chat.completions.create(
                 model=self.model_name,
                 messages=messages,
-                temperature=0.0,
+                temperature=temperature,
                 max_tokens=max_tokens,
                 top_p=1,
                 frequency_penalty=0,
@@ -384,97 +375,3 @@ class OpenAIGenerator:
             tokens_used += response.usage.total_tokens
             contents.append(content)
         return contents, tokens_used
-
-
-class OAMGenerator:
-    def __init__(
-        self,
-        model_name,
-    ):
-        self.model_name = model_name
-        logger.info("Loading model")
-        time_now = time.time()
-        self.model = PlancraftOAM.from_pretrained(model_name)
-        self.model.cuda()
-        logger.info(f"Model loaded in {time.time() - time_now:.2f} seconds")
-        # compile
-        time_now = time.time()
-        self.model = torch.compile(self.model)
-        logger.info(f"Model compiled in {time.time() - time_now:.2f} seconds")
-
-        self.model.eval()
-        self.model.tokenizer.pad_token = self.model.tokenizer.eos_token
-        self.model.tokenizer.truncation_side = "left"
-
-    def reset(self):
-        pass
-
-    def prepare_messages(
-        self,
-        history: History,
-        max_messages_window: int,
-        system_prompt: dict = None,
-        prompt_images: list = [],
-    ) -> tuple[list[dict], list]:
-        """
-        Prepare the messages using a history
-        """
-        message_window = history.dialogue_history[-max_messages_window:]
-        # remove the first assistant message if it is present
-        if len(message_window) > 0 and message_window[0]["role"] == "assistant":
-            message_window = message_window[1:]
-        # add the system prompt if the first message is not a system message
-        if message_window[0]["role"] != "system" and system_prompt is not None:
-            message_window = [system_prompt] + message_window
-
-        image_window = []
-
-        image_list = prompt_images + history.images
-        image_count = 0
-
-        # iterate through the messages to count how many images are present
-        new_message_window = []
-        for m in message_window:
-            message_content = m["content"]
-            message_role = m["role"]
-
-            if "\ninventory:\n" in message_content:
-                message_content = (
-                    message_content.split("\ninventory:\n")[0]
-                    + "\ninventory:<|inventory|>"
-                )
-
-            if "<|inventory|>" in message_content:
-                image_count += 1
-
-            new_message_window.append(
-                {"role": message_role, "content": message_content}
-            )
-
-        assert image_count <= len(image_list), "Too many images"
-        # get messages from end of queue
-        image_window = image_list[-image_count:]
-        image_window = [Image.fromarray(img) for img in image_window]
-
-        return new_message_window, image_window
-
-    @torch.inference_mode()
-    def generate_unconstrained(
-        self,
-        batch_messages: list[list[dict]],
-        images: list,
-        max_tokens: int = 256,
-        temperature=0.6,
-        **kwargs,
-    ) -> tuple[list[str], int]:
-        """
-        Generate unconstrained text based on the batch of messages.
-        """
-        text_responses, total_tokens_used = self.model.generate(
-            batch_messages=batch_messages,
-            batch_images=images,
-            do_sample=True,
-            temperature=temperature,
-            max_new_tokens=max_tokens,
-        )
-        return text_responses, total_tokens_used

plancraft/models/oracle.py

@@ -1,59 +1,51 @@
 import copy
 from collections import Counter
 
+import torch
+
 from plancraft.config import EvalConfig
-from plancraft.environments.actions import (
+from plancraft.environment.actions import (
+    MoveAction,
+    SmeltAction,
     StopAction,
-    SymbolicMoveAction,
-    SymbolicSmeltAction,
 )
-from plancraft.environments.planner import optimal_planner
-from plancraft.environments.recipes import (
+from plancraft.environment.planner import optimal_planner
+from plancraft.environment.recipes import (
     ShapedRecipe,
     ShapelessRecipe,
     SmeltingRecipe,
     id_to_item,
 )
-from plancraft.environments.sampler import MAX_STACK_SIZE
-from plancraft.models.base import ABCModel, History
+from plancraft.environment.sampler import MAX_STACK_SIZE
+from plancraft.models.base import PlancraftBaseModel
+from plancraft.models.bbox_model import IntegratedBoundingBoxModel
 
 
 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:
+def find_free_inventory_slot(inventory: 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
+    from_item_type = inventory[from_slot]["type"]
+    from_item_quantity = inventory[from_slot]["quantity"]
+
+    empty_slots = set(range(10, 46)) - set(inventory.keys()) - set([from_slot])
 
     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
+    for slot, item in inventory.items():
+        if slot == from_slot:
+            continue
         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["quantity"] == 0:
+        #     item_type = "air"
         if item_type not in type_to_slot:
-            type_to_slot[item_type] = [item_slot]
+            type_to_slot[item_type] = [slot]
         else:
-            type_to_slot[item_type].append(item_slot)
+            type_to_slot[item_type].append(slot)
 
-        if item_slot not in slot_to_quantity:
-            slot_to_quantity[item_slot] = quantity
-        else:
-            slot_to_quantity[item_slot] += quantity
+        slot_to_quantity[slot] = item["quantity"]
 
     assert from_item_type is not None, f"Item not found in slot {from_slot}"
 
@@ -61,111 +53,108 @@ def find_free_inventory_slot(inventory: list[dict], from_slot: int) -> int:
     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
+                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
+    # if "air" in type_to_slot:
+    #     for slot in type_to_slot["air"]:
+    #         if slot > 10:
+    #             return slot
+
+    if len(empty_slots) > 0:
+        return empty_slots.pop()
 
     raise ValueError("No free slot found")
 
 
-def find_item_in_inventory(target: str, inventory: list[dict]) -> int:
-    for item in inventory:
+def find_item_in_inventory(target: str, inventory: dict) -> int:
+    for slot, item in inventory.items():
         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")
+            return slot
 
 
-def get_inventory_counter(inventory: list[dict]) -> Counter:
+def get_inventory_counter(inventory: 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":
+    for slot, item in inventory.items():
+        if slot == 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:
+def get_crafting_slot_item(inventory: dict) -> dict:
+    for slot, item in inventory.items():
+        if slot == 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]:
+    inventory: dict, slot_from: int, slot_to: int, quantity: int
+) -> 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)
+    new_inventory = dict(inventory)
+    new_inventory[slot_from]["quantity"] -= quantity
     return new_inventory
 
 
-class OracleModel(ABCModel):
+class OracleModel(PlancraftBaseModel):
     """
     Oracle model returns actions that solve the task optimally
     """
 
     def __init__(self, cfg: EvalConfig):
-        self._history = History(objective="")
         self.plans = []
         self.subplans = []
+        self.use_fasterrcnn = cfg.plancraft.use_fasterrcnn
 
-    @property
-    def history(self):
-        return self._history
+        self.bbox_model = None
+        if self.use_fasterrcnn:
+            # fasterrcnn is not multimodal model but a separate model
+            self.bbox_model = IntegratedBoundingBoxModel.from_pretrained(
+                "gautierdag/plancraft-fasterrcnn"
+            )
+            self.bbox_model.eval()
+            if torch.cuda.is_available():
+                self.bbox_model.cuda()
 
-    def reset_history(self, objective: str = ""):
-        self.history.reset(objective=objective)
+    def reset(self):
         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)
+        return optimal_planner(
+            target=observation["target"], inventory=inventory_counter
+        )
 
-    def get_next_action(
-        self, observation: dict
-    ) -> SymbolicMoveAction | SymbolicSmeltAction:
+    def get_next_action(self, observation: dict) -> MoveAction | SmeltAction:
         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"])
+        if self.bbox_model is not None:
+            observed_inventory = self.bbox_model.get_inventory(
+                observation["image"].copy()
+            )
+        else:
+            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(
+            return MoveAction(
                 slot_from=0, slot_to=free_slot, quantity=slot_item["quantity"]
             )
 
@@ -176,7 +165,8 @@ class OracleModel(ABCModel):
         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 not self.use_fasterrcnn:
+            assert len(new_items) == 1
 
         if isinstance(plan_recipe, ShapelessRecipe):
             crafting_slot = 1
@@ -193,8 +183,7 @@ class OracleModel(ABCModel):
                         n += 1
                         continue
 
-                    # low_level_plan.append(("move", item, from_slot, crafting_slot, 1))
-                    action = SymbolicMoveAction(
+                    action = MoveAction(
                         slot_from=from_slot, slot_to=crafting_slot, quantity=1
                     )
                     # update state of inventory
@@ -202,7 +191,6 @@ class OracleModel(ABCModel):
                         current_inventory, from_slot, crafting_slot, 1
                     )
                     self.subplans.append(action)
-
                     crafting_slot += 1
                     n += 1
 
@@ -214,7 +202,7 @@ class OracleModel(ABCModel):
                 free_slot = find_free_inventory_slot(
                     current_inventory, from_slot=from_slot
                 )
-                action = SymbolicSmeltAction(
+                action = SmeltAction(
                     slot_from=from_slot, slot_to=free_slot, quantity=quantity
                 )
                 self.subplans.append(action)
@@ -228,7 +216,7 @@ class OracleModel(ABCModel):
                     for item in valid_items:
                         if items_to_use_counter[item] > 0:
                             from_slot = find_item_in_inventory(item, current_inventory)
-                            action = SymbolicMoveAction(
+                            action = MoveAction(
                                 slot_from=from_slot,
                                 slot_to=inventory_position,
                                 quantity=1,
@@ -245,19 +233,13 @@ class OracleModel(ABCModel):
 
         return self.subplans.pop(0)
 
-    def step(self, observation: dict) -> list[SymbolicMoveAction | SymbolicSmeltAction]:
-        # add observation to history
-        self.history.add_observation_to_history(observation)
-
+    def step(self, observation: dict, **kwargs) -> str:
         # get action
         if len(self.plans) == 0:
-            self.get_plan(observation)
+            self.plans = self.get_plan(observation)
             if self.plans is None:
                 self.plans = []
-                return StopAction()
+                return str(StopAction())
 
         action = self.get_next_action(observation)
-
-        # add action to history
-        self.history.add_action_to_history(action)
-        return action
+        return str(action)