cc

scripts/generate_dataset.sh

@@ -1,13 +1,13 @@
-mkdir -p data
-mkdir -p data/apigen
-mkdir -p data/glaive
-mkdir -p data/toolace
+mkdir -p datasets
+mkdir -p datasets/apigen
+mkdir -p datasets/glaive
+mkdir -p datasets/toolace
 
 # Download datasets
-# wget -O data/apigen/xlam_function_calling_60k.json https://huggingface.co/datasets/Salesforce/xlam-function-calling-60k/resolve/main/xlam_function_calling_60k.json
+# wget -O datasets/apigen/xlam_function_calling_60k.json https://huggingface.co/datasets/Salesforce/xlam-function-calling-60k/resolve/main/xlam_function_calling_60k.json
 echo "APIGen cannot be downloaded directly due to permission issues. Please download it manually and place it in the datasets/apigen folder."
-wget -O data/glaive/glaive-function-calling-v2.json https://huggingface.co/datasets/glaiveai/glaive-function-calling-v2/resolve/main/glaive-function-calling-v2.json
-wget -O data/toolace/data.json https://huggingface.co/datasets/Team-ACE/ToolACE/resolve/main/data.json
+wget -O datasets/glaive/glaive-function-calling-v2.json https://huggingface.co/datasets/glaiveai/glaive-function-calling-v2/resolve/main/glaive-function-calling-v2.json
+wget -O datasets/toolace/data.json https://huggingface.co/datasets/Team-ACE/ToolACE/resolve/main/data.json
 
 # Preprocess datasets
 python ./scripts/preprocess_apigen_dataset.py

scripts/mix_datasets.py

@@ -6,7 +6,7 @@ DATASETS = ["apigen", "glaive", "toolace"]
 data = []
 
 for dataset_name in DATASETS:
-    with open(f"./data/{dataset_name}/output.json") as f:
+    with open(f"./datasets/{dataset_name}/output.json") as f:
         subdata = json.load(f)
 
     subdata = [{**item, "source": dataset_name} for item in subdata]
@@ -86,6 +86,12 @@ from random import shuffle, seed
 one_two = [item for item in data if len(item["used_tools"]) in [1, 2]]
 other = [item for item in data if len(item["used_tools"]) > 2]
 
+seed(42)
+shuffle(one_two)
+
+seed(42)
+shuffle(other)
+
 train_samples = one_two[: int(len(one_two) * 0.8)]
 test_samples = one_two[int(len(one_two) * 0.8) :] + other
 
@@ -118,10 +124,10 @@ print("Number of samples in test dataset:", len(test_dataset["samples"]))
 
 import os
 
-os.makedirs("./data/mixed", exist_ok=True)
+os.makedirs("./datasets/mixed", exist_ok=True)
 
-with open("./data/mixed/train.json", "w") as f:
+with open("./datasets/mixed/train.json", "w") as f:
     json.dump(train_dataset, f, indent=2)
 
-with open("./data/mixed/test.json", "w") as f:
+with open("./datasets/mixed/test.json", "w") as f:
     json.dump(test_dataset, f, indent=2)

scripts/preprocess_apigen_dataset.py

@@ -1,6 +1,6 @@
 import json
 
-with open("./data/apigen/xlam_function_calling_60k.json", "r") as f:
+with open("./datasets/apigen/xlam_function_calling_60k.json", "r") as f:
     data = json.load(f)
 
 results = []
@@ -23,5 +23,5 @@ for sample in tqdm(data, desc="Processing APIGen samples"):
 
     results.append(result)
 
-with open("./data/apigen/output.json", "w") as f:
+with open("./datasets/apigen/output.json", "w") as f:
     json.dump(results, f, indent=4)

scripts/preprocess_glaive_dataset.py

@@ -2,7 +2,7 @@ import re
 import json
 
 
-with open("./data/glaive/glaive-function-calling-v2.json", "r") as f:
+with open("./datasets/glaive/glaive-function-calling-v2.json", "r") as f:
     data = json.load(f)
 
 
@@ -53,5 +53,5 @@ for sample in tqdm(data, desc="Processing GLAIVE samples"):
     if processed:
         results.extend(processed)
 
-with open("./data/glaive/output.json", "w") as f:
+with open("./datasets/glaive/output.json", "w") as f:
     json.dump(results, f, indent=2)

scripts/preprocess_toolace_dataset.py

@@ -1,7 +1,7 @@
 import json
 import re
 
-with open("./data/toolace/data.json") as f:
+with open("./datasets/toolace/data.json") as f:
     data = json.load(f)
 
 
@@ -58,5 +58,5 @@ for sample in tqdm(data, desc="Processing ToolACE samples"):
     if processed:
         results.extend(processed)
 
-with open("./data/toolace/output.json", "w") as f:
+with open("./datasets/toolace/output.json", "w") as f:
     json.dump(results, f, indent=2)

src/models/miniagent_module.py

@@ -67,10 +67,6 @@ class MiniAgentModule(LightningModule):
         pred = self.pred_model(inst_emb_r, tool_emb_r)  # [BxB, 1]
         pred = pred.view(B, B)  # [B, B]
 
-        # 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
 
@@ -146,5 +142,5 @@ class MiniAgentModule(LightningModule):
         pass
 
     def configure_optimizers(self):
-        opt = torch.optim.AdamW(self.parameters(), lr=self.lr)
+        opt = torch.optim.AdamW(self.parameters(), lr=self.lr, weight_decay=1e-4)
         return opt

src/models/mlp_module.py

@@ -7,6 +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.linear2 = nn.Linear(hidden_dim, output_dim)
 
     def forward(self, x_output):
@@ -15,6 +16,7 @@ class MLPProjection(nn.Module):
 
         x = self.linear1(x)
         x = F.silu(x)
+        x = self.dropout(x)
         x = self.linear2(x)
 
         return x
@@ -30,13 +32,15 @@ class MLPPrediction(nn.Module):
         real_input_dim = input_dim * (2 + int(use_abs_diff) + int(use_mult))
 
         self.mlp = nn.Sequential(
-            nn.Linear(real_input_dim, 1024),
-            nn.SiLU(),
-            nn.Linear(1024, 512),
+            nn.Linear(real_input_dim, 512),
             nn.SiLU(),
+            nn.Dropout(0.3),
             nn.Linear(512, 256),
             nn.SiLU(),
-            nn.Linear(256, 1),
+            nn.Dropout(0.3),
+            nn.Linear(256, 128),
+            nn.SiLU(),
+            nn.Linear(128, 1),
         )
 
     def forward(self, x1, x2):

test_bert.ipynb

@@ -154,7 +154,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "datamodule = MixedDataModule(dataset_path=\"./data/mixed\", batch_size=32, num_workers=4, bert_model=\"bert-base-uncased\", tool_capacity=16)"
+    "datamodule = MixedDataModule(dataset_path=\"./datasets/mixed\", batch_size=32, num_workers=4, bert_model=\"bert-base-uncased\", tool_capacity=16)"
    ]
   },
   {