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

train/dataset.py

@@ -1,187 +0,0 @@
-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