cc

src/data/mixed_datamodule.py

@@ -6,7 +6,7 @@ from .mixed_dataset import MixedDataset
 
 class MixedDataModule(LightningDataModule):
     def __init__(
-        self, bert_model, dataset_path, tool_capacity, batch_size, num_workers
+        self, bert_model, dataset_path, tool_capacity, batch_size, num_workers, seed
     ):
         super().__init__()
         self.bert_model = bert_model
@@ -14,6 +14,7 @@ class MixedDataModule(LightningDataModule):
         self.tool_capacity = tool_capacity
         self.batch_size = batch_size
         self.num_workers = num_workers
+        self.seed = seed
 
     def setup(self, stage=None):
         if stage == "fit":
@@ -22,12 +23,14 @@ class MixedDataModule(LightningDataModule):
                 "train",
                 os.path.join(self.dataset_path, "train.json"),
                 self.tool_capacity,
+                seed=self.seed,
             )
             self.val_dataset = MixedDataset(
                 self.bert_model,
                 "test",
                 os.path.join(self.dataset_path, "test.json"),
                 self.tool_capacity,
+                seed=self.seed,
             )
         elif stage == "test":
             self.test_dataset = MixedDataset(
@@ -35,6 +38,7 @@ class MixedDataModule(LightningDataModule):
                 "test",
                 os.path.join(self.dataset_path, "test.json"),
                 self.tool_capacity,
+                seed=self.seed,
             )
 
     def train_dataloader(self):

src/data/mixed_dataset.py

@@ -26,7 +26,7 @@ class MixedDataset(Dataset):
         return tools, samples
 
     def encode_text(self, text):
-        inputs = self.tokenizer.encode_plus(
+        inputs = self.tokenizer(
             text,
             max_length=128,
             padding="max_length",
@@ -53,8 +53,8 @@ class MixedDataset(Dataset):
             return {
                 "inst_ids": inst_ids,
                 "inst_mask": inst_mask,
-                "tool_desc_ids": tool_desc_ids,
-                "tool_desc_mask": tool_desc_mask,
+                "tool_ids": tool_desc_ids,
+                "tool_mask": tool_desc_mask,
             }
         else:
             # for testing, we sample a random set of tools + the correct tool, size = tool_capacity
@@ -74,7 +74,7 @@ class MixedDataset(Dataset):
             )
 
             tools = correct_tools + wrong_tools
-            tool_ids, tool_ids_mask = self.encode_text(
+            tool_ids, tool_mask = self.encode_text(
                 [self.tools[tool_id]["description"] for tool_id in tools]
             )
 
@@ -82,6 +82,6 @@ class MixedDataset(Dataset):
                 "inst_ids": inst_ids,
                 "inst_mask": inst_mask,
                 "tool_ids": tool_ids,
-                "tool_ids_mask": tool_ids_mask,
+                "tool_mask": tool_mask,
                 "correct_tool_mask": correct_tool_mask,
             }

src/models/miniagent_module.py

@@ -33,7 +33,7 @@ class MiniAgentModule(LightningModule):
         self.tool_proj_model = tool_proj_model
         self.pred_model = pred_model
 
-        self.val_acc = Accuracy()
+        self.val_acc = Accuracy(task="binary")
 
         self.lr = lr
 
@@ -50,8 +50,8 @@ class MiniAgentModule(LightningModule):
 
         inst_ids = batch["inst_ids"]
         inst_mask = batch["inst_mask"]
-        tool_ids = batch["tool_desc_ids"]
-        tool_mask = batch["tool_desc_mask"]
+        tool_ids = batch["tool_ids"]
+        tool_mask = batch["tool_mask"]
 
         inst_z = self.bert_model(inst_ids, inst_mask, return_dict=False)[1]
         tool_z = self.bert_model(tool_ids, tool_mask, return_dict=False)[1]
@@ -82,8 +82,8 @@ class MiniAgentModule(LightningModule):
     ) -> None:
         inst_ids = batch["inst_ids"]
         inst_mask = batch["inst_mask"]
-        tool_ids = batch["tool_desc_ids"]
-        tool_mask = batch["tool_desc_mask"]
+        tool_ids = batch["tool_ids"]
+        tool_mask = batch["tool_mask"]
         correct_tool_mask = batch["correct_tool_mask"]
 
         B = inst_ids.shape[0]  # batch size