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

plancraft/train/dataset.py

@@ -0,0 +1,187 @@
+import glob
+import json
+import random
+
+import torch
+from torch.utils.data import Dataset
+
+TEMPLATES = {
+    "idefics2": {
+        "assistant": "\nAssistant:",
+        "user": "\nUser:",
+    },
+    "llama3": {
+        "assistant": "<|start_header_id|>assistant<|end_header_id|>\n\n",
+        "user": "<|start_header_id|>user<|end_header_id|>\n\n",
+    },
+}
+
+
+class PlancraftDialogueDataset(Dataset):
+    def __init__(
+        self,
+        dataset_dir: str = "data/oracle",
+        use_images=False,
+        trace_mode="oa",
+        split="train",
+        max_message_window=30,
+    ):
+        super().__init__()
+        self.split = split
+        self.use_images = use_images
+        self.trace_mode = trace_mode
+        self.add_system_message = True
+
+        assert trace_mode in ["oa", "ota"], f"Invalid trace mode {trace_mode}"
+
+        print("Loading dialogue dataset")
+        data = []
+        for example_path in sorted(
+            glob.glob(f"{dataset_dir}/{split}/{trace_mode}/*.json")
+        ):
+            with open(example_path) as f:
+                example = json.load(f)
+            data.append(example)
+        self.dataset = data
+        self.max_message_window = max_message_window
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def __getitem__(self, idx: int) -> tuple[dict, list]:
+        example = self.dataset[idx]
+        if len(example["messages"]) > self.max_message_window:
+            # add system message
+            if self.add_system_message:
+                messages = [example["messages"][0]]
+            else:
+                messages = []
+
+            # sample window
+            user_messages_idxs = list(
+                range(self.max_message_window, len(example["messages"]), 2)
+            )
+            end = random.choice(user_messages_idxs)
+            start = end - self.max_message_window + 1
+            assert start != 0
+            # add window
+            messages = messages + example["messages"][start : end + 1]
+
+        else:
+            messages = example["messages"]
+        return messages
+
+
+def track_assistant_response(
+    batch,
+    tokenizer,
+    template_name: str = "llama3",
+):
+    """
+    Mask that returns 1 for tokens in the assistant response and 0 otherwise.
+    """
+    assistant_template = TEMPLATES[template_name]["assistant"]
+    user_template = TEMPLATES[template_name]["user"]
+    start_seq = tokenizer.encode(
+        assistant_template,
+        add_special_tokens=False,
+        return_tensors="pt",
+    )[0]
+    end_seq = tokenizer.encode(
+        user_template,
+        add_special_tokens=False,
+        return_tensors="pt",
+    )[0]
+    encoded_label_ids = batch["labels"]
+    mask = torch.zeros_like(encoded_label_ids)
+    for seq_idx, seq in enumerate(encoded_label_ids):
+        in_masked_response = False
+        i = 0
+        while i < len(seq):
+            if i + len(start_seq) < len(seq) and torch.all(
+                seq[i : i + len(start_seq)].eq(start_seq)
+            ):
+                in_masked_response = True
+                i += len(start_seq)
+                continue
+            if i + len(end_seq) < len(seq) and torch.all(
+                seq[i : i + len(end_seq)].eq(end_seq)
+            ):
+                in_masked_response = False
+                i += len(end_seq)
+                continue
+            if in_masked_response:
+                mask[seq_idx, i] = 1
+            else:
+                mask[seq_idx, i] = 0
+            i += 1
+    return mask
+
+
+def get_collate_fn(
+    tokenizer,
+    max_length=8142,
+    only_assistant=False,
+    template_name: str = "llama3",
+):
+    assert template_name in TEMPLATES
+
+    def collate_fn(batch):
+        messages_batch = []
+        for messages in batch:
+            text = tokenizer.apply_chat_template(
+                messages, add_generation_prompt=False, tokenize=False
+            )
+            # remove BOS token since it will later be added again by the tokenizer
+            text = text.replace("<|begin_of_text|>", "")
+            messages_batch.append(text)
+
+        batch = tokenizer(
+            messages_batch,
+            truncation=True,
+            max_length=max_length,
+            return_tensors="pt",
+        )
+        labels = batch["input_ids"].clone()
+        batch["labels"] = labels
+
+        # add mask for assistant response
+        if only_assistant:
+            mask = track_assistant_response(
+                batch, tokenizer, template_name=template_name
+            )
+            labels[mask == 0] = -100
+
+        return batch
+
+    return collate_fn
+
+
+def get_dataset_and_collate(
+    tokenizer,
+    template_name: str = "llama3",
+    max_length: int = 8142,
+    max_message_window: int = 30,
+    trace_mode="oa",
+    only_assistant=False,
+):
+    if template_name == "llama3":
+        train_dataset = PlancraftDialogueDataset(
+            use_images=False,
+            max_message_window=max_message_window,
+            split="train",
+            trace_mode=trace_mode,
+        )
+        val_dataset = PlancraftDialogueDataset(
+            use_images=False,
+            max_message_window=max_message_window,
+            split="val",
+            trace_mode=trace_mode,
+        )
+        collate_fn = get_collate_fn(
+            tokenizer=tokenizer,
+            only_assistant=only_assistant,
+            template_name=template_name,
+            max_length=max_length,
+        )
+    return train_dataset, val_dataset, collate_fn

plancraft/utils.py

@@ -0,0 +1,84 @@
+import glob
+import pathlib
+
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+
+def get_downloaded_models() -> dict:
+    """
+    Get the list of downloaded models on the NFS partition (EIDF).
+    """
+    downloaded_models = {}
+    # known models on NFS partition
+    if pathlib.Path("/nfs").exists():
+        local_models = glob.glob("/nfs/public/hf/models/*/*")
+        downloaded_models = {
+            model.replace("/nfs/public/hf/models/", ""): model for model in local_models
+        }
+    return downloaded_models
+
+
+def get_torch_device() -> torch.device:
+    device = torch.device("cpu")
+    if torch.cuda.is_available():
+        device = torch.device("cuda", 0)
+    elif torch.backends.mps.is_available():
+        if not torch.backends.mps.is_built():
+            print(
+                "MPS not available because the current PyTorch install was not built with MPS enabled."
+            )
+        else:
+            device = torch.device("mps")
+    return device
+
+
+def resize_image(img, target_resolution=(128, 128)):
+    if type(img) == np.ndarray:
+        img = cv2.resize(img, target_resolution, interpolation=cv2.INTER_LINEAR)
+    elif type(img) == torch.Tensor:
+        img = F.interpolate(img, size=target_resolution, mode="bilinear")
+    else:
+        raise ValueError("Unsupported type for img")
+    return img
+
+
+def save_frames_to_video(frames: list, out_path: str):
+    imgs = []
+    for id, (frame, goal) in enumerate(frames):
+        # if torch.is_tensor(frame):
+        # frame = frame.permute(0, 2, 3, 1).cpu().numpy()
+
+        frame = resize_image(frame, (320, 240)).astype("uint8")
+        cv2.putText(
+            frame,
+            f"FID: {id}",
+            (10, 25),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.8,
+            (255, 255, 255),
+            2,
+        )
+        cv2.putText(
+            frame,
+            f"Goal: {goal}",
+            (10, 55),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.8,
+            (255, 0, 0),
+            2,
+        )
+        imgs.append(Image.fromarray(frame))
+    imgs = imgs[::3]
+    imgs[0].save(
+        out_path,
+        save_all=True,
+        append_images=imgs[1:],
+        optimize=False,
+        quality=0,
+        duration=150,
+        loop=0,
+    )

{plancraft-0.1.2.dist-info → plancraft-0.1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: plancraft
-Version: 0.1.2
+Version: 0.1.4
 Summary: Plancraft: an evaluation dataset for planning with LLM agents
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown