cc

configs/callbacks/default.yaml

@@ -8,7 +8,7 @@ defaults:
 model_checkpoint:
   dirpath: ${paths.output_dir}/checkpoints
   filename: "epoch_{epoch:03d}"
-  monitor: "val/acc"
+  monitor: "val/1_acc"
   mode: "max"
   save_last: True
   auto_insert_metric_name: False

src/models/miniagent_module.py

@@ -31,9 +31,17 @@ class MiniAgentModule(LightningModule):
         self.tool_proj_model = tool_proj_model
         self.pred_model = pred_model
 
-        self.val_acc = Accuracy(task="binary")
-        self.val_precision = MeanMetric()
-        self.val_recall = MeanMetric()
+        self.val_1_acc = Accuracy(task="binary")
+        self.val_1_precision = MeanMetric()
+        self.val_1_recall = MeanMetric()
+
+        self.val_2_acc = Accuracy(task="binary")
+        self.val_2_precision = MeanMetric()
+        self.val_2_recall = MeanMetric()
+
+        self.val_other_acc = Accuracy(task="binary")
+        self.val_other_precision = MeanMetric()
+        self.val_other_recall = MeanMetric()
 
         self.lr = lr
 
@@ -69,8 +77,6 @@ class MiniAgentModule(LightningModule):
 
         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)
 
@@ -107,31 +113,151 @@ class MiniAgentModule(LightningModule):
         pred = torch.sigmoid(pred)
 
         pred_tool_mask = pred > 0.5
-        pos_sample = (pred_tool_mask == correct_tool_mask).all(dim=1).long()
 
         true_pos_mask = pred_tool_mask & correct_tool_mask
 
-        precision = true_pos_mask.sum(dim=1) / torch.clamp(
-            pred_tool_mask.sum(dim=1), min=1
+        one_tool_mask = correct_tool_mask.sum(dim=1) == 1
+        two_tool_mask = correct_tool_mask.sum(dim=1) == 2
+        other_mask = ~(one_tool_mask | two_tool_mask)
+
+        # one tool
+        one_tool_pos_sample = (
+            (pred_tool_mask[one_tool_mask] == correct_tool_mask[one_tool_mask])
+            .all(dim=1)
+            .long()
+        )
+
+        one_tool_precision = true_pos_mask[one_tool_mask].sum(dim=1) / torch.clamp(
+            pred_tool_mask[one_tool_mask].sum(dim=1), min=1
+        )
+
+        one_tool_recall = true_pos_mask[one_tool_mask].sum(dim=1) / torch.clamp(
+            correct_tool_mask[one_tool_mask].sum(dim=1), min=1
+        )
+
+        # two tool
+        two_tool_pos_sample = (
+            (pred_tool_mask[two_tool_mask] == correct_tool_mask[two_tool_mask])
+            .all(dim=1)
+            .long()
         )
-        recall = true_pos_mask.sum(dim=1) / correct_tool_mask.sum(dim=1)
-
-        self.val_acc.update(pos_sample, torch.ones_like(pos_sample))
-        self.val_precision.update(precision)
-        self.val_recall.update(recall)
-
-        self.log("val/acc", self.val_acc, on_epoch=True, sync_dist=True, prog_bar=True)
-        self.log(
-            "val/precision",
-            self.val_precision,
-            on_epoch=True,
-            sync_dist=True,
-            prog_bar=True,
+
+        two_tool_precision = true_pos_mask[two_tool_mask].sum(dim=1) / torch.clamp(
+            pred_tool_mask[two_tool_mask].sum(dim=1), min=1
         )
-        self.log(
-            "val/recall", self.val_recall, on_epoch=True, sync_dist=True, prog_bar=True
+
+        two_tool_recall = true_pos_mask[two_tool_mask].sum(dim=1) / torch.clamp(
+            correct_tool_mask[two_tool_mask].sum(dim=1), min=1
         )
 
+        # other
+        other_pos_sample = (
+            (pred_tool_mask[other_mask] == correct_tool_mask[other_mask])
+            .all(dim=1)
+            .long()
+        )
+
+        other_precision = true_pos_mask[other_mask].sum(dim=1) / torch.clamp(
+            pred_tool_mask[other_mask].sum(dim=1), min=1
+        )
+
+        other_recall = true_pos_mask[other_mask].sum(dim=1) / torch.clamp(
+            correct_tool_mask[other_mask].sum(dim=1), min=1
+        )
+
+        if one_tool_pos_sample.sum().item() > 0:
+            self.val_1_acc.update(
+                one_tool_pos_sample, torch.ones_like(one_tool_pos_sample)
+            )
+            self.val_1_precision.update(one_tool_precision)
+            self.val_1_recall.update(one_tool_recall)
+
+            self.log(
+                "val/1_acc",
+                self.val_1_acc,
+                on_epoch=True,
+                sync_dist=True,
+                prog_bar=True,
+            )
+
+            self.log(
+                "val/1_precision",
+                self.val_1_precision,
+                on_epoch=True,
+                sync_dist=True,
+                prog_bar=True,
+            )
+
+            self.log(
+                "val/1_recall",
+                self.val_1_recall,
+                on_epoch=True,
+                sync_dist=True,
+                prog_bar=True,
+            )
+
+        if two_tool_pos_sample.sum().item() > 0:
+            self.val_2_acc.update(
+                two_tool_pos_sample, torch.ones_like(two_tool_pos_sample)
+            )
+            self.val_2_precision.update(two_tool_precision)
+            self.val_2_recall.update(two_tool_recall)
+
+            self.log(
+                "val/2_acc",
+                self.val_2_acc,
+                on_epoch=True,
+                sync_dist=True,
+                prog_bar=True,
+            )
+
+            self.log(
+                "val/2_precision",
+                self.val_2_precision,
+                on_epoch=True,
+                sync_dist=True,
+                prog_bar=True,
+            )
+
+            self.log(
+                "val/2_recall",
+                self.val_2_recall,
+                on_epoch=True,
+                sync_dist=True,
+                prog_bar=True,
+            )
+
+        if other_pos_sample.sum().item() > 0:
+            self.val_other_acc.update(
+                other_pos_sample, torch.ones_like(other_pos_sample)
+            )
+            self.val_other_precision.update(other_precision)
+            self.val_other_recall.update(other_recall)
+
+            self.log(
+                "val/other_acc",
+                self.val_other_acc,
+                on_epoch=True,
+                sync_dist=True,
+                prog_bar=True,
+            )
+
+            self.log(
+                "val/other_precision",
+                self.val_other_precision,
+                on_epoch=True,
+                sync_dist=True,
+                prog_bar=True,
+            )
+
+            self.log(
+                "val/other_recall",
+                self.val_other_recall,
+                on_epoch=True,
+                sync_dist=True,
+                prog_bar=True,
+            )
+
     def on_validation_epoch_end(self) -> None:
         pass
 

src/models/mlp_module.py

@@ -7,7 +7,7 @@ class MLPProjection(nn.Module):
     def __init__(self, input_dim, hidden_dim, output_dim):
         super().__init__()
         self.linear1 = nn.Linear(input_dim, hidden_dim)
-        self.dropout = nn.Dropout(0.3)
+        self.dropout = nn.Dropout(0.5)
         self.linear2 = nn.Linear(hidden_dim, output_dim)
 
     def forward(self, x_output):
@@ -34,10 +34,10 @@ class MLPPrediction(nn.Module):
         self.mlp = nn.Sequential(
             nn.Linear(real_input_dim, 512),
             nn.SiLU(),
-            nn.Dropout(0.3),
+            nn.Dropout(0.5),
             nn.Linear(512, 256),
             nn.SiLU(),
-            nn.Dropout(0.3),
+            nn.Dropout(0.5),
             nn.Linear(256, 128),
             nn.SiLU(),
             nn.Linear(128, 1),