cc

configs/experiment/miniagent-bert-mlp-abs_diff.yaml

@@ -0,0 +1,37 @@
+# @package _global_
+
+defaults:
+  - override /data: mixed
+  - override /model: miniagent
+  - override /callbacks: default
+  - override /trainer: gpu
+
+seed: 42
+
+model:
+  lr: 0.001
+  bert_model: bert-base-uncased
+
+  inst_proj_model:
+    _target_: src.models.mlp_module.MLPProjection
+    input_dim: 768
+    hidden_dim: 768
+    output_dim: 768
+
+  tool_proj_model:
+    _target_: src.models.mlp_module.MLPProjection
+    input_dim: 768
+    hidden_dim: 768
+    output_dim: 768
+
+  pred_model:
+    _target_: src.models.mlp_module.MLPPrediction
+    input_dim: 768
+    use_abs_diff: true
+    use_mult: true
+
+data:
+  bert_model: bert-base-uncased
+  seed: 42
+  batch_size: 128
+  tool_capacity: 16

configs/experiment/miniagent-bert-mlp-mult.yaml

@@ -0,0 +1,37 @@
+# @package _global_
+
+defaults:
+  - override /data: mixed
+  - override /model: miniagent
+  - override /callbacks: default
+  - override /trainer: gpu
+
+seed: 42
+
+model:
+  lr: 0.001
+  bert_model: bert-base-uncased
+
+  inst_proj_model:
+    _target_: src.models.mlp_module.MLPProjection
+    input_dim: 768
+    hidden_dim: 768
+    output_dim: 768
+
+  tool_proj_model:
+    _target_: src.models.mlp_module.MLPProjection
+    input_dim: 768
+    hidden_dim: 768
+    output_dim: 768
+
+  pred_model:
+    _target_: src.models.mlp_module.MLPPrediction
+    input_dim: 768
+    use_abs_diff: true
+    use_mult: true
+
+data:
+  bert_model: bert-base-uncased
+  seed: 42
+  batch_size: 128
+  tool_capacity: 16

configs/trainer/default.yaml

@@ -3,7 +3,7 @@ _target_: lightning.pytorch.trainer.Trainer
 default_root_dir: ${paths.output_dir}
 
 min_epochs: 1 # prevents early stopping
-max_epochs: 20
+max_epochs: 50
 
 accelerator: cpu
 devices: 1
@@ -18,4 +18,4 @@ check_val_every_n_epoch: 1
 
 # set True to to ensure deterministic results
 # makes training slower but gives more reproducibility than just setting seeds
-deterministic: False
+deterministic: True

src/data/mixed_dataset.py

@@ -25,13 +25,20 @@ class MixedDataset(Dataset):
 
         return tools, samples
 
-    def encode_text(self, text):
-        inputs = self.tokenizer(
-            text,
-            max_length=128,
-            padding="max_length",
-            truncation=True,
-        )
+    def encode_text(self, text, padding=True):
+        if padding:
+            inputs = self.tokenizer(
+                text,
+                max_length=128,
+                padding="max_length",
+                truncation=True,
+            )
+        else:
+            inputs = self.tokenizer(
+                text,
+                max_length=128,
+                truncation=True,
+            )
         ids = torch.tensor(inputs["input_ids"], dtype=torch.long)
         mask = torch.tensor(inputs["attention_mask"], dtype=torch.long)
 

src/models/miniagent_module.py

@@ -26,8 +26,8 @@ class MiniAgentModule(LightningModule):
         )
 
         self.bert_model = BertModel.from_pretrained(bert_model)
-        self.bert_model.eval()
-        self.bert_model.requires_grad_(False)
+        # self.bert_model.eval()
+        # self.bert_model.requires_grad_(False)
 
         self.inst_proj_model = inst_proj_model
         self.tool_proj_model = tool_proj_model
@@ -67,8 +67,12 @@ class MiniAgentModule(LightningModule):
         pred = self.pred_model(inst_emb_r, tool_emb_r)  # [BxB, 1]
         pred = pred.view(B, B)  # [B, B]
 
-        labels = torch.arange(B, device=pred.device).long()
-        loss = (F.cross_entropy(pred, labels) + F.cross_entropy(pred.T, labels)) * 0.5
+        target = torch.eye(B, device=pred.device).float()
+
+        pos_weight = torch.tensor([B - 1], device=pred.device)
+        loss = F.binary_cross_entropy_with_logits(pred, target, pos_weight=pos_weight)
+        # labels = torch.arange(B, device=pred.device).long()
+        # loss = (F.cross_entropy(pred, labels) + F.cross_entropy(pred.T, labels)) * 0.5
 
         self.log("train/loss", loss, on_step=True, sync_dist=True, prog_bar=True)
 
@@ -142,5 +146,16 @@ class MiniAgentModule(LightningModule):
         pass
 
     def configure_optimizers(self):
-        opt = torch.optim.AdamW(self.parameters(), lr=self.lr, weight_decay=1e-4)
+        opt = torch.optim.AdamW(
+            [
+                {"params": self.bert_model.parameters(), "lr": 1e-5},
+                {
+                    "params": list(self.inst_proj_model.parameters())
+                    + list(self.tool_proj_model.parameters())
+                    + list(self.pred_model.parameters()),
+                    "lr": self.lr,
+                },
+            ],
+            weight_decay=1e-4,
+        )
         return opt