cc

.gitignore

@@ -1,4 +1,4 @@
-data
+datasets
 logs
 wandb
 __pycache__

src/data/mixed_datamodule.py

@@ -0,0 +1,62 @@
+import os
+from torch.utils.data import DataLoader
+from lightning import LightningDataModule
+from .mixed_dataset import MixedDataset
+
+
+class MixedDataModule(LightningDataModule):
+    def __init__(
+        self, bert_model, dataset_path, tool_capacity, batch_size, num_workers
+    ):
+        super().__init__()
+        self.bert_model = bert_model
+        self.dataset_path = dataset_path
+        self.tool_capacity = tool_capacity
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+
+    def setup(self, stage=None):
+        if stage == "fit":
+            self.train_dataset = MixedDataset(
+                self.bert_model,
+                "train",
+                os.path.join(self.dataset_path, "train.json"),
+                self.tool_capacity,
+            )
+            self.val_dataset = MixedDataset(
+                self.bert_model,
+                "test",
+                os.path.join(self.dataset_path, "test.json"),
+                self.tool_capacity,
+            )
+        elif stage == "test":
+            self.test_dataset = MixedDataset(
+                self.bert_model,
+                "test",
+                os.path.join(self.dataset_path, "test.json"),
+                self.tool_capacity,
+            )
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+        )
+
+    def test_dataloader(self):
+        return DataLoader(
+            self.test_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+        )

src/data/mixed_dataset.py

@@ -0,0 +1,76 @@
+import os
+import json
+import torch
+import random
+from torch.utils.data import Dataset
+from transformers import BertTokenizer
+
+
+class MixedDataset(Dataset):
+    def __init__(self, bert_model, stage, anno_file, tool_capacity):
+        self.stage = stage
+        self.tool_capacity = tool_capacity
+        self.tools, self.samples = self.load_data(anno_file)
+        self.tool_ids = list(self.tools.keys())
+        self.tokenizer = BertTokenizer.from_pretrained(bert_model)
+
+    def load_data(self, anno_file):
+        with open(anno_file, "r") as f:
+            data = json.load(f)
+        tools = data["tools"]
+        samples = data["samples"]
+
+        tools = {tool["id"]: tool for tool in tools}
+
+        return tools, samples
+
+    def encode_text(self, text):
+        inputs = self.tokenizer.encode_plus(
+            text,
+            max_length=128,
+            padding="max_length",
+            truncation=True,
+        )
+        ids = torch.tensor(inputs["input_ids"], dtype=torch.long)
+        mask = torch.tensor(inputs["attention_mask"], dtype=torch.long)
+
+        return ids, mask
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        sample = self.samples[idx]
+        inst = sample["instruction"]
+        inst_ids, inst_mask = self.encode_text(inst)
+
+        if self.stage == "train":
+            tool_id = random.choice(sample["tools"])
+            tool_desc = self.tools[tool_id]["description"]
+            tool_desc_ids, tool_desc_mask = self.encode_text(tool_desc)
+
+            return {
+                "inst_ids": inst_ids,
+                "inst_mask": inst_mask,
+                "tool_desc_ids": tool_desc_ids,
+                "tool_desc_mask": tool_desc_mask,
+            }
+        else:
+            # for testing, we sample a random set of tools + the correct tool, size = tool_capacity
+            # wrong tools are sampled randomly from self.tools
+            correct_tools = sample["tools"]
+            wrong_tools = random.sample(
+                [tool for tool in self.tool_ids if tool not in correct_tools],
+                self.tool_capacity - len(correct_tools),
+            )
+            tools = correct_tools + wrong_tools
+            tool_ids, tool_ids_mask = self.encode_text(
+                [self.tools[tool_id]["description"] for tool_id in tools]
+            )
+
+            return {
+                "inst_ids": inst_ids,
+                "inst_mask": inst_mask,
+                "tool_ids": tool_ids,
+                "tool_ids_mask": tool_ids_mask,
+            }