cc

configs/experiment/example.yaml

@@ -1,41 +0,0 @@
-# @package _global_
-
-# to execute this experiment run:
-# python train.py experiment=example
-
-defaults:
-  - override /data: mnist
-  - override /model: mnist
-  - override /callbacks: default
-  - override /trainer: default
-
-# all parameters below will be merged with parameters from default configurations set above
-# this allows you to overwrite only specified parameters
-
-tags: ["mnist", "simple_dense_net"]
-
-seed: 12345
-
-trainer:
-  min_epochs: 10
-  max_epochs: 10
-  gradient_clip_val: 0.5
-
-model:
-  optimizer:
-    lr: 0.002
-  net:
-    lin1_size: 128
-    lin2_size: 256
-    lin3_size: 64
-  compile: false
-
-data:
-  batch_size: 64
-
-logger:
-  wandb:
-    tags: ${tags}
-    group: "mnist"
-  aim:
-    experiment: "mnist"

configs/experiment/miniagent-bert-mlp.yaml

@@ -0,0 +1,36 @@
+# @package _global_
+
+defaults:
+  - override /data: mixed
+  - override /model: miniagent
+  - override /callbacks: default
+  - override /trainer: gpu
+
+seed: 42
+
+model:
+  lr: 0.0001
+  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: false
+    use_mult: false
+
+data:
+  bert_model: bert-base-uncased
+  batch_size: 32
+  tool_capacity: 16

configs/model/miniagent.yaml

@@ -0,0 +1,3 @@
+_target_: src.models.miniagent_module.MiniAgentModule
+
+bert_model: bert-base-uncased

configs/model/mnist.yaml

@@ -1,25 +0,0 @@
-_target_: src.models.mnist_module.MNISTLitModule
-
-optimizer:
-  _target_: torch.optim.Adam
-  _partial_: true
-  lr: 0.001
-  weight_decay: 0.0
-
-scheduler:
-  _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
-  _partial_: true
-  mode: min
-  factor: 0.1
-  patience: 10
-
-net:
-  _target_: src.models.components.simple_dense_net.SimpleDenseNet
-  input_size: 784
-  lin1_size: 64
-  lin2_size: 128
-  lin3_size: 64
-  output_size: 10
-
-# compile model for faster training with pytorch 2.0
-compile: false

configs/train.yaml

@@ -4,11 +4,11 @@
 # order of defaults determines the order in which configs override each other
 defaults:
   - _self_
-  - data: mnist
-  - model: mnist
+  - data: mixed
+  - model: miniagent
   - callbacks: default
-  - logger: null # set logger here or use command line (e.g. `python train.py logger=tensorboard`)
-  - trainer: default
+  - logger: wandb # set logger here or use command line (e.g. `python train.py logger=tensorboard`)
+  - trainer: gpu
   - paths: default
   - extras: default
   - hydra: default
@@ -46,4 +46,4 @@ test: True
 ckpt_path: null
 
 # seed for random number generators in pytorch, numpy and python.random
-seed: null
+seed: 42

notebooks/test_bert.ipynb

@@ -1,244 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import torch\n",
-    "from transformers import BertModel, BertTokenizer\n",
-    "\n",
-    "model_name = \"bert-base-uncased\"\n",
-    "\n",
-    "tokenizer = BertTokenizer.from_pretrained(model_name)\n",
-    "model = BertModel.from_pretrained(model_name).cuda()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 19,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "tensor([[  101,  2182,  2003,  2070,  3793,  2000,  4372, 16044,   102]],\n",
-      "       device='cuda:0')\n"
-     ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "BaseModelOutputWithPoolingAndCrossAttentions(last_hidden_state=tensor([[[-0.0549,  0.1053, -0.1065,  ..., -0.3550,  0.0686,  0.6506],\n",
-       "         [-0.5759, -0.3650, -0.1383,  ..., -0.6782,  0.2092, -0.1639],\n",
-       "         [-0.1641, -0.5597,  0.0150,  ..., -0.1603, -0.1346,  0.6216],\n",
-       "         ...,\n",
-       "         [ 0.2448,  0.1254,  0.1587,  ..., -0.2749, -0.1163,  0.8809],\n",
-       "         [ 0.0481,  0.4950, -0.2827,  ..., -0.6097, -0.1212,  0.2527],\n",
-       "         [ 0.9046,  0.2137, -0.5897,  ...,  0.3040, -0.6172, -0.1950]]],\n",
-       "       device='cuda:0'), pooler_output=tensor([[-8.3875e-01, -2.8128e-01, -4.9307e-02,  5.9384e-01,  1.9040e-01,\n",
-       "         -1.3149e-01,  8.4180e-01,  1.8583e-01, -3.4094e-03, -9.9988e-01,\n",
-       "          2.5326e-02,  3.7435e-01,  9.6628e-01, -2.6636e-01,  8.5380e-01,\n",
-       "         -4.6390e-01,  4.3777e-02, -5.0858e-01,  2.4810e-01, -5.1466e-01,\n",
-       "          4.9031e-01,  9.9296e-01,  5.5225e-01,  2.2509e-01,  3.0543e-01,\n",
-       "          5.2188e-01, -4.6495e-01,  8.9979e-01,  9.2244e-01,  6.7778e-01,\n",
-       "         -5.8705e-01,  1.7724e-01, -9.7322e-01, -1.3662e-01, -1.8714e-01,\n",
-       "         -9.7217e-01,  2.5207e-01, -7.0827e-01,  2.4030e-04,  2.0537e-02,\n",
-       "         -8.6379e-01,  3.1709e-01,  9.9879e-01, -3.7660e-01,  7.7044e-02,\n",
-       "         -2.8340e-01, -9.9993e-01,  1.8332e-01, -8.3515e-01,  1.1565e-01,\n",
-       "          5.0053e-02, -4.2741e-02,  1.2181e-01,  3.1478e-01,  3.5848e-01,\n",
-       "          2.1106e-01, -1.7044e-01,  1.2083e-01, -8.6361e-02, -5.1711e-01,\n",
-       "         -4.9923e-01,  2.6499e-01, -2.1260e-01, -8.7598e-01,  7.0440e-03,\n",
-       "         -2.9079e-01, -1.7843e-02, -2.3522e-01, -1.9458e-02, -8.6378e-02,\n",
-       "          8.0144e-01,  1.1972e-01,  1.0071e-01, -7.2592e-01, -1.2743e-01,\n",
-       "          9.6182e-02, -4.8009e-01,  1.0000e+00, -2.8258e-01, -9.5353e-01,\n",
-       "          1.4182e-01,  6.0877e-02,  4.2072e-01,  5.1125e-01, -3.7690e-01,\n",
-       "         -1.0000e+00,  3.7079e-01, -1.1285e-01, -9.7725e-01,  5.0917e-02,\n",
-       "          4.0374e-01, -8.9133e-02, -2.4458e-01,  3.9134e-01, -1.7817e-01,\n",
-       "         -1.8376e-01, -2.1744e-01, -6.2693e-02, -1.6466e-01, -1.1514e-01,\n",
-       "          7.3975e-02, -1.4703e-01, -2.8427e-02, -3.2398e-01,  1.3225e-01,\n",
-       "         -2.4507e-01, -3.8186e-01,  2.8680e-01, -2.2497e-01,  5.3901e-01,\n",
-       "          3.1464e-01, -2.2506e-01,  1.5128e-01, -9.1770e-01,  5.3941e-01,\n",
-       "         -2.4810e-01, -9.6917e-01, -4.1745e-01, -9.7295e-01,  5.8796e-01,\n",
-       "          8.2848e-02, -5.6085e-02,  9.2621e-01,  4.6422e-01,  2.3537e-01,\n",
-       "          5.4746e-02, -2.8870e-01, -1.0000e+00, -2.8243e-01, -3.2437e-01,\n",
-       "          1.5666e-01, -7.7765e-02, -9.6036e-01, -9.2680e-01,  5.7584e-01,\n",
-       "          9.4303e-01,  9.6574e-02,  9.9754e-01, -1.0749e-01,  8.8828e-01,\n",
-       "          1.8555e-01, -1.5084e-01, -2.7866e-01, -3.2091e-01,  5.0980e-01,\n",
-       "          3.2259e-01, -6.5085e-01,  1.1357e-01,  2.4492e-02,  2.6022e-02,\n",
-       "         -2.6347e-01, -2.0440e-01, -3.4154e-02, -8.8905e-01, -3.0404e-01,\n",
-       "          8.9583e-01,  5.9022e-02, -1.0150e-02,  6.4437e-01, -1.0694e-01,\n",
-       "         -3.0653e-01,  7.8505e-01,  4.7535e-01,  2.1861e-01,  1.8127e-02,\n",
-       "          3.4923e-01, -8.8995e-02,  4.1078e-01, -8.3798e-01,  1.8740e-01,\n",
-       "          3.0741e-01, -1.6962e-01, -9.0649e-03, -9.5766e-01, -2.1796e-01,\n",
-       "          4.8940e-01,  9.7536e-01,  7.1374e-01,  2.6859e-01,  2.7887e-01,\n",
-       "         -2.5057e-01,  2.8507e-01, -9.0920e-01,  9.5826e-01, -4.7828e-02,\n",
-       "          1.8272e-01,  2.4859e-01, -1.3071e-02, -8.0613e-01, -4.1807e-01,\n",
-       "          7.6876e-01, -2.1496e-01, -7.9618e-01,  6.0786e-02, -4.1543e-01,\n",
-       "         -3.0691e-01, -1.7626e-01,  4.7289e-01, -2.2226e-01, -3.4807e-01,\n",
-       "          1.9591e-02,  8.9078e-01,  9.3557e-01,  6.5839e-01, -4.2014e-01,\n",
-       "          4.0678e-01, -8.4855e-01, -3.4254e-01,  4.9145e-02,  1.6904e-01,\n",
-       "          4.2381e-02,  9.8512e-01, -3.0301e-01, -1.8377e-01, -8.8705e-01,\n",
-       "         -9.6226e-01, -7.3838e-02, -8.2263e-01, -1.6584e-02, -6.1115e-01,\n",
-       "          3.1010e-01,  2.1122e-01, -5.6624e-02,  3.1030e-01, -9.6320e-01,\n",
-       "         -7.0888e-01,  2.6261e-01, -2.6109e-01,  3.0822e-01, -2.0598e-01,\n",
-       "          6.7753e-01,  1.7591e-01, -5.3198e-01,  5.1929e-01,  8.7353e-01,\n",
-       "          1.1513e-02, -6.4033e-01,  7.1928e-01, -2.0344e-01,  8.7310e-01,\n",
-       "         -5.2567e-01,  9.6434e-01,  3.6636e-01,  3.8012e-01, -8.5828e-01,\n",
-       "          6.6280e-02, -8.5531e-01,  1.3883e-01, -1.5715e-01, -5.6898e-01,\n",
-       "          7.5188e-02,  4.8760e-01,  2.2753e-01,  7.3960e-01, -4.1109e-01,\n",
-       "          9.8629e-01, -5.8852e-01, -9.2256e-01, -1.0970e-01, -4.9473e-03,\n",
-       "         -9.6941e-01,  2.2550e-01,  1.4062e-01, -4.1525e-02, -3.5449e-01,\n",
-       "         -4.7126e-01, -9.2133e-01,  8.0581e-01,  4.2704e-02,  9.7325e-01,\n",
-       "          1.5150e-02, -8.2761e-01, -2.6887e-01, -8.6111e-01, -2.1065e-01,\n",
-       "         -5.0521e-02,  4.8313e-01, -2.3626e-01, -8.9746e-01,  3.6525e-01,\n",
-       "          5.3635e-01,  3.2864e-01,  1.7254e-01,  9.8927e-01,  9.9986e-01,\n",
-       "          9.5408e-01,  7.9837e-01,  8.2591e-01, -9.7353e-01, -3.6506e-01,\n",
-       "          9.9994e-01, -6.9189e-01, -9.9999e-01, -9.0511e-01, -4.9764e-01,\n",
-       "          3.5348e-01, -1.0000e+00, -7.4901e-02,  4.5044e-02, -8.3963e-01,\n",
-       "         -2.2185e-01,  9.6232e-01,  9.6733e-01, -1.0000e+00,  8.2459e-01,\n",
-       "          9.1231e-01, -4.8469e-01,  3.5679e-01, -2.6093e-01,  9.5376e-01,\n",
-       "          3.0112e-01,  4.6869e-01, -6.6399e-02,  2.9968e-01, -3.6107e-01,\n",
-       "         -7.6258e-01,  1.2195e-01,  1.8867e-01,  8.9893e-01,  7.6925e-02,\n",
-       "         -6.9103e-01, -8.5666e-01,  6.2071e-02,  5.7222e-02, -3.0562e-01,\n",
-       "         -9.2684e-01, -1.4870e-01, -2.9001e-01,  5.3266e-01,  6.7110e-02,\n",
-       "          1.9160e-01, -7.2418e-01,  2.3380e-01, -5.4946e-01,  3.2379e-01,\n",
-       "          5.2877e-01, -9.0674e-01, -4.9937e-01, -4.4697e-02, -4.0249e-01,\n",
-       "          1.7510e-01, -9.4833e-01,  9.4143e-01, -2.0001e-01,  3.7983e-01,\n",
-       "          1.0000e+00,  4.7085e-02, -7.9982e-01,  3.6463e-01,  1.4701e-01,\n",
-       "         -4.6405e-01,  1.0000e+00,  4.2325e-01, -9.6105e-01, -3.5190e-01,\n",
-       "          3.0839e-01, -3.3907e-01, -3.5049e-01,  9.9734e-01, -8.8495e-02,\n",
-       "          1.4489e-01,  2.2866e-01,  9.4882e-01, -9.7813e-01,  8.1165e-01,\n",
-       "         -8.5137e-01, -9.2814e-01,  9.3252e-01,  8.8071e-01, -1.3548e-01,\n",
-       "         -6.6040e-01, -1.6448e-02, -1.4088e-01,  1.9752e-01, -9.2127e-01,\n",
-       "          4.6474e-01,  2.7750e-01, -2.0809e-02,  8.2319e-01, -7.3596e-01,\n",
-       "         -4.2808e-01,  3.5810e-01,  3.2605e-02,  2.9415e-01,  3.4245e-01,\n",
-       "          3.8726e-01, -1.9187e-01, -2.6267e-02, -1.0283e-01, -2.9876e-01,\n",
-       "         -9.5093e-01,  1.9958e-01,  1.0000e+00,  1.1309e-01, -6.9028e-02,\n",
-       "         -9.7448e-02,  1.6147e-02, -2.6026e-01,  3.4048e-01,  4.3058e-01,\n",
-       "         -2.0093e-01, -7.9416e-01,  8.8086e-02, -8.1918e-01, -9.6890e-01,\n",
-       "          5.9739e-01,  1.2771e-01, -2.6104e-01,  9.9899e-01,  1.4233e-01,\n",
-       "          1.2059e-01,  4.0483e-02,  5.1663e-01, -2.8824e-02,  4.7300e-01,\n",
-       "         -2.7801e-01,  9.5372e-01, -2.0937e-01,  4.0244e-01,  7.6304e-01,\n",
-       "          3.1611e-02, -1.7255e-01, -5.7354e-01,  3.6617e-03, -8.8957e-01,\n",
-       "          2.3939e-02, -9.0446e-01,  9.4090e-01,  2.0094e-02,  2.2362e-01,\n",
-       "          6.5433e-02,  7.9535e-02,  1.0000e+00, -3.0244e-01,  5.7777e-01,\n",
-       "         -1.9970e-01,  7.3194e-01, -9.7151e-01, -6.7480e-01, -3.0641e-01,\n",
-       "          4.5864e-02,  2.2067e-01, -1.9735e-01,  1.2545e-01, -9.4334e-01,\n",
-       "          2.7122e-04, -1.1515e-01, -9.5089e-01, -9.7541e-01,  4.2593e-01,\n",
-       "          6.3330e-01, -5.9799e-02, -7.2403e-01, -5.1824e-01, -5.4331e-01,\n",
-       "          3.2069e-01, -1.8098e-01, -8.6600e-01,  4.9465e-01, -2.0962e-01,\n",
-       "          3.1231e-01, -1.9074e-01,  4.5443e-01, -7.5650e-02,  8.2584e-01,\n",
-       "         -1.6397e-01, -1.7857e-02, -5.9298e-02, -7.1197e-01,  6.4738e-01,\n",
-       "         -7.2854e-01, -2.4250e-01, -8.7492e-02,  1.0000e+00, -3.0222e-01,\n",
-       "          1.1153e-01,  6.3838e-01,  5.6054e-01, -9.7328e-02,  1.4136e-01,\n",
-       "          2.4002e-01,  7.8386e-02,  4.0799e-01,  1.1524e-01, -4.4807e-01,\n",
-       "         -1.6638e-01,  4.3632e-01, -1.1570e-01, -1.5105e-01,  6.8226e-01,\n",
-       "          5.2354e-01,  1.7305e-03,  2.1510e-02,  4.1700e-03,  9.9460e-01,\n",
-       "          1.7238e-02, -1.2763e-01, -3.8213e-01,  2.8577e-02, -2.2359e-01,\n",
-       "         -3.2380e-01,  1.0000e+00,  2.7486e-01,  2.9497e-02, -9.7830e-01,\n",
-       "         -1.3164e-01, -8.6022e-01,  9.9958e-01,  6.8541e-01, -7.7604e-01,\n",
-       "          4.9727e-01,  2.5391e-01, -1.3586e-01,  5.6488e-01, -1.2485e-01,\n",
-       "         -2.2596e-01,  8.4730e-02,  9.8264e-03,  9.0287e-01, -3.3443e-01,\n",
-       "         -9.3577e-01, -5.7664e-01,  3.0516e-01, -9.2533e-01,  9.8399e-01,\n",
-       "         -5.3663e-01, -1.1600e-01, -2.4484e-01,  1.1554e-01,  4.8479e-01,\n",
-       "         -1.2667e-01, -9.4730e-01, -1.2690e-01, -8.1969e-03,  9.2896e-01,\n",
-       "          8.2588e-02, -3.9092e-01, -8.7582e-01, -1.1210e-01,  1.4961e-01,\n",
-       "          3.9152e-02, -8.8999e-01,  9.4467e-01, -9.5203e-01,  2.1941e-01,\n",
-       "          9.9999e-01,  2.9063e-01, -4.6173e-01,  8.3155e-02, -2.9862e-01,\n",
-       "          2.1412e-01,  1.2559e-01,  4.1466e-01, -9.2596e-01, -1.5259e-01,\n",
-       "         -1.1740e-01,  1.7836e-01, -7.1510e-02,  6.5254e-02,  5.5836e-01,\n",
-       "          1.7951e-01, -3.6205e-01, -4.6712e-01, -6.8167e-02,  2.7963e-01,\n",
-       "          6.4896e-01, -1.6497e-01, -1.5638e-01,  1.0768e-01, -1.3088e-01,\n",
-       "         -7.8675e-01, -1.7540e-01, -2.6264e-01, -9.9611e-01,  5.5635e-01,\n",
-       "         -1.0000e+00, -1.6508e-01, -4.3545e-01, -1.5236e-01,  7.6126e-01,\n",
-       "          3.3181e-01,  1.6770e-01, -6.2815e-01,  2.1790e-01,  7.3190e-01,\n",
-       "          6.5308e-01, -1.4397e-01,  9.8390e-02, -6.1118e-01,  1.8836e-01,\n",
-       "         -9.3982e-02,  2.3265e-01,  1.1375e-01,  7.3512e-01, -8.6529e-02,\n",
-       "          1.0000e+00,  4.7590e-02, -4.0572e-01, -9.1862e-01,  2.3070e-01,\n",
-       "         -1.1442e-01,  9.9996e-01, -8.1991e-01, -9.1503e-01,  2.0697e-01,\n",
-       "         -5.1682e-01, -7.5368e-01,  2.1045e-01,  5.6595e-03, -6.4855e-01,\n",
-       "         -4.1527e-01,  9.1653e-01,  7.3400e-01, -4.2874e-01,  2.3967e-01,\n",
-       "         -2.4068e-01, -3.3491e-01, -4.3406e-02, -2.8789e-02,  9.7527e-01,\n",
-       "          3.9052e-01,  8.5009e-01,  4.9386e-01, -7.6392e-03,  9.4010e-01,\n",
-       "          2.0048e-01,  2.7198e-01,  6.1171e-02,  1.0000e+00,  2.8885e-01,\n",
-       "         -9.1301e-01,  2.4983e-01, -9.7476e-01, -1.2421e-01, -9.2382e-01,\n",
-       "          1.9579e-01,  8.3662e-02,  8.1386e-01, -2.0425e-01,  9.2494e-01,\n",
-       "          5.8940e-02,  2.2477e-02,  1.0070e-01,  3.0503e-01,  2.4897e-01,\n",
-       "         -8.5819e-01, -9.6919e-01, -9.6845e-01,  3.7154e-01, -3.5970e-01,\n",
-       "         -1.5696e-02,  2.3336e-01,  6.0670e-02,  3.0880e-01,  2.7858e-01,\n",
-       "         -1.0000e+00,  8.9189e-01,  2.8790e-01, -9.3187e-02,  9.3518e-01,\n",
-       "          2.2574e-01,  3.3557e-01,  1.7759e-01, -9.6978e-01, -9.1210e-01,\n",
-       "         -2.6055e-01, -3.0031e-01,  7.1668e-01,  5.6109e-01,  7.4522e-01,\n",
-       "          2.2312e-01, -4.3589e-01, -3.5411e-01,  2.3049e-01, -4.5362e-01,\n",
-       "         -9.8112e-01,  3.6383e-01,  7.7364e-02, -9.1412e-01,  9.2768e-01,\n",
-       "         -4.3448e-01, -7.7045e-02,  5.9378e-01, -1.9775e-01,  8.6886e-01,\n",
-       "          6.9895e-01,  3.6932e-01,  3.9024e-02,  4.8328e-01,  7.9945e-01,\n",
-       "          9.2656e-01,  9.7223e-01, -6.6530e-02,  6.7691e-01,  1.3522e-01,\n",
-       "          3.0958e-01,  5.4615e-01, -9.1295e-01,  8.0810e-02,  1.6158e-01,\n",
-       "          6.6348e-02,  1.2711e-01, -1.4847e-01, -9.2298e-01,  4.7369e-01,\n",
-       "         -1.1713e-01,  4.0759e-01, -3.2167e-01,  1.7547e-01, -3.2665e-01,\n",
-       "         -1.0958e-01, -7.2304e-01, -3.8936e-01,  4.9849e-01,  1.8104e-01,\n",
-       "          8.7080e-01,  4.7525e-01,  5.3199e-02, -5.8020e-01, -7.1312e-02,\n",
-       "          2.1172e-01, -8.8360e-01,  8.8993e-01,  9.6810e-02,  4.6417e-01,\n",
-       "          8.0528e-02, -1.0677e-01,  7.8678e-01, -2.3742e-01, -3.0270e-01,\n",
-       "         -1.4502e-01, -7.0351e-01,  7.5302e-01, -3.6792e-01, -4.0656e-01,\n",
-       "         -2.9551e-01,  5.9936e-01,  2.2145e-01,  9.9465e-01,  3.4475e-02,\n",
-       "         -5.7931e-02, -2.5533e-01, -2.1419e-01,  3.0497e-01, -2.0330e-01,\n",
-       "         -1.0000e+00,  3.2221e-01, -2.2759e-02, -1.8658e-02, -9.2918e-02,\n",
-       "          2.2858e-02, -1.0043e-01, -9.5813e-01, -1.5660e-01,  2.8550e-01,\n",
-       "         -2.3666e-02, -4.3055e-01, -1.8355e-01,  3.7819e-01,  4.1632e-01,\n",
-       "          5.1126e-01,  8.4217e-01,  3.9612e-02,  4.3960e-01,  4.4244e-01,\n",
-       "         -2.3278e-01, -5.7333e-01,  8.3078e-01]], device='cuda:0'), hidden_states=None, past_key_values=None, attentions=None, cross_attentions=None)"
-      ]
-     },
-     "execution_count": 19,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "input_text = \"Here is some text to encode\"\n",
-    "input_ids = tokenizer.encode(input_text, add_special_tokens=True)\n",
-    "input_ids = torch.tensor([input_ids]).cuda()\n",
-    "\n",
-    "print(input_ids)\n",
-    "\n",
-    "with torch.no_grad():\n",
-    "    last_hidden_states = model(input_ids)  # Models outputs are now tuples\n",
-    "\n",
-    "last_hidden_states"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "swim",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.12.6"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}

src/eval.py

@@ -1,5 +1,6 @@
 from typing import Any, Dict, List, Tuple
 
+import torch
 import hydra
 import rootutils
 from lightning import LightningDataModule, LightningModule, Trainer
@@ -34,6 +35,8 @@ from src.utils import (
 
 log = RankedLogger(__name__, rank_zero_only=True)
 
+torch.set_float32_matmul_precision("medium")
+
 
 @task_wrapper
 def evaluate(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:

src/models/components/simple_dense_net.py

@@ -1,54 +0,0 @@
-import torch
-from torch import nn
-
-
-class SimpleDenseNet(nn.Module):
-    """A simple fully-connected neural net for computing predictions."""
-
-    def __init__(
-        self,
-        input_size: int = 784,
-        lin1_size: int = 256,
-        lin2_size: int = 256,
-        lin3_size: int = 256,
-        output_size: int = 10,
-    ) -> None:
-        """Initialize a `SimpleDenseNet` module.
-
-        :param input_size: The number of input features.
-        :param lin1_size: The number of output features of the first linear layer.
-        :param lin2_size: The number of output features of the second linear layer.
-        :param lin3_size: The number of output features of the third linear layer.
-        :param output_size: The number of output features of the final linear layer.
-        """
-        super().__init__()
-
-        self.model = nn.Sequential(
-            nn.Linear(input_size, lin1_size),
-            nn.BatchNorm1d(lin1_size),
-            nn.ReLU(),
-            nn.Linear(lin1_size, lin2_size),
-            nn.BatchNorm1d(lin2_size),
-            nn.ReLU(),
-            nn.Linear(lin2_size, lin3_size),
-            nn.BatchNorm1d(lin3_size),
-            nn.ReLU(),
-            nn.Linear(lin3_size, output_size),
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Perform a single forward pass through the network.
-
-        :param x: The input tensor.
-        :return: A tensor of predictions.
-        """
-        batch_size, channels, width, height = x.size()
-
-        # (batch, 1, width, height) -> (batch, 1*width*height)
-        x = x.view(batch_size, -1)
-
-        return self.model(x)
-
-
-if __name__ == "__main__":
-    _ = SimpleDenseNet()

src/models/miniagent_module.py

@@ -0,0 +1,96 @@
+from typing import Any, Dict, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from lightning import LightningModule
+from torchmetrics import MaxMetric, MeanMetric
+from torchmetrics.classification.accuracy import Accuracy
+
+from transformers import BertModel
+
+
+class MiniAgentModule(LightningModule):
+    def __init__(
+        self,
+        bert_model: str,
+        inst_proj_model: nn.Module,
+        tool_proj_model: nn.Module,
+        pred_model: nn.Module,
+        lr: float,
+    ) -> None:
+        super().__init__()
+
+        self.save_hyperparameters(
+            logger=False, ignore=["inst_proj_model", "tool_proj_model", "pred_model"]
+        )
+
+        self.bert_model = BertModel.from_pretrained(bert_model)
+        self.bert_model.eval()
+        self.bert_model.requires_grad_(False)
+
+        self.inst_proj_model = inst_proj_model
+        self.tool_proj_model = tool_proj_model
+        self.pred_model = pred_model
+
+        self.lr = lr
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        pass
+
+    def on_train_start(self) -> None:
+        pass
+
+    def training_step(
+        self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int
+    ) -> torch.Tensor:
+        B = batch["inst_ids"].shape[0]
+
+        inst_ids = batch["inst_ids"]
+        inst_mask = batch["inst_mask"]
+        tool_ids = batch["tool_desc_ids"]
+        tool_mask = batch["tool_desc_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]
+
+        inst_emb = self.inst_proj_model(inst_z)
+        tool_emb = self.tool_proj_model(tool_z)
+
+        inst_emb_r = inst_emb.unsqueeze(0).repeat(B, 1, 1).view(B * B, -1)
+        tool_emb_r = tool_emb.unsqueeze(1).repeat(1, B, 1).view(B * B, -1)
+
+        pred = self.pred_model(inst_emb_r, tool_emb_r)  # [BxB, 1]
+        pred = pred.view(B, B)  # [B, B]
+
+        # mask out the diagonal
+        target = torch.eye(B, device=pred.device).float()
+
+        loss = F.binary_cross_entropy_with_logits(pred, target)
+
+        self.log("train/loss", loss, on_step=True, sync_dist=True, prog_bar=True)
+
+        return loss
+
+    def on_train_epoch_end(self) -> None:
+        pass
+
+    def validation_step(
+        self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int
+    ) -> None:
+        pass
+
+    def on_validation_epoch_end(self) -> None:
+        pass
+
+    def test_step(
+        self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int
+    ) -> None:
+        pass
+
+    def on_test_epoch_end(self) -> None:
+        pass
+
+    def configure_optimizers(self):
+        opt = torch.optim.AdamW(self.parameters(), lr=self.lr)
+        return opt

src/models/mlp_module.py

@@ -0,0 +1,55 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MLPProjection(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim):
+        super().__init__()
+        self.linear1 = nn.Linear(input_dim, hidden_dim)
+        self.linear2 = nn.Linear(hidden_dim, output_dim)
+
+    def forward(self, x_output):
+        # only use first token ([CLS]) of each output
+        x = x_output
+
+        x = self.linear1(x)
+        x = F.silu(x)
+        x = self.linear2(x)
+
+        return x
+
+
+class MLPPrediction(nn.Module):
+    def __init__(self, input_dim, use_abs_diff=False, use_mult=False):
+        super().__init__()
+
+        self.use_abs_diff = use_abs_diff
+        self.use_mult = use_mult
+
+        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.SiLU(),
+            nn.Linear(512, 256),
+            nn.SiLU(),
+            nn.Linear(256, 1),
+        )
+
+    def forward(self, x1, x2):
+        x = torch.cat([x1, x2], dim=1)
+
+        if self.use_abs_diff:
+            x_diff = torch.abs(x1 - x2)
+            x = torch.cat([x, x_diff], dim=1)
+
+        if self.use_mult:
+            x_mult = x1 * x2
+            x = torch.cat([x, x_mult], dim=1)
+
+        x = self.mlp(x)
+
+        return x

src/models/mnist_module.py

@@ -1,217 +0,0 @@
-from typing import Any, Dict, Tuple
-
-import torch
-from lightning import LightningModule
-from torchmetrics import MaxMetric, MeanMetric
-from torchmetrics.classification.accuracy import Accuracy
-
-
-class MNISTLitModule(LightningModule):
-    """Example of a `LightningModule` for MNIST classification.
-
-    A `LightningModule` implements 8 key methods:
-
-    ```python
-    def __init__(self):
-    # Define initialization code here.
-
-    def setup(self, stage):
-    # Things to setup before each stage, 'fit', 'validate', 'test', 'predict'.
-    # This hook is called on every process when using DDP.
-
-    def training_step(self, batch, batch_idx):
-    # The complete training step.
-
-    def validation_step(self, batch, batch_idx):
-    # The complete validation step.
-
-    def test_step(self, batch, batch_idx):
-    # The complete test step.
-
-    def predict_step(self, batch, batch_idx):
-    # The complete predict step.
-
-    def configure_optimizers(self):
-    # Define and configure optimizers and LR schedulers.
-    ```
-
-    Docs:
-        https://lightning.ai/docs/pytorch/latest/common/lightning_module.html
-    """
-
-    def __init__(
-        self,
-        net: torch.nn.Module,
-        optimizer: torch.optim.Optimizer,
-        scheduler: torch.optim.lr_scheduler,
-        compile: bool,
-    ) -> None:
-        """Initialize a `MNISTLitModule`.
-
-        :param net: The model to train.
-        :param optimizer: The optimizer to use for training.
-        :param scheduler: The learning rate scheduler to use for training.
-        """
-        super().__init__()
-
-        # this line allows to access init params with 'self.hparams' attribute
-        # also ensures init params will be stored in ckpt
-        self.save_hyperparameters(logger=False)
-
-        self.net = net
-
-        # loss function
-        self.criterion = torch.nn.CrossEntropyLoss()
-
-        # metric objects for calculating and averaging accuracy across batches
-        self.train_acc = Accuracy(task="multiclass", num_classes=10)
-        self.val_acc = Accuracy(task="multiclass", num_classes=10)
-        self.test_acc = Accuracy(task="multiclass", num_classes=10)
-
-        # for averaging loss across batches
-        self.train_loss = MeanMetric()
-        self.val_loss = MeanMetric()
-        self.test_loss = MeanMetric()
-
-        # for tracking best so far validation accuracy
-        self.val_acc_best = MaxMetric()
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Perform a forward pass through the model `self.net`.
-
-        :param x: A tensor of images.
-        :return: A tensor of logits.
-        """
-        return self.net(x)
-
-    def on_train_start(self) -> None:
-        """Lightning hook that is called when training begins."""
-        # by default lightning executes validation step sanity checks before training starts,
-        # so it's worth to make sure validation metrics don't store results from these checks
-        self.val_loss.reset()
-        self.val_acc.reset()
-        self.val_acc_best.reset()
-
-    def model_step(
-        self, batch: Tuple[torch.Tensor, torch.Tensor]
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Perform a single model step on a batch of data.
-
-        :param batch: A batch of data (a tuple) containing the input tensor of images and target labels.
-
-        :return: A tuple containing (in order):
-            - A tensor of losses.
-            - A tensor of predictions.
-            - A tensor of target labels.
-        """
-        x, y = batch
-        logits = self.forward(x)
-        loss = self.criterion(logits, y)
-        preds = torch.argmax(logits, dim=1)
-        return loss, preds, y
-
-    def training_step(
-        self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int
-    ) -> torch.Tensor:
-        """Perform a single training step on a batch of data from the training set.
-
-        :param batch: A batch of data (a tuple) containing the input tensor of images and target
-            labels.
-        :param batch_idx: The index of the current batch.
-        :return: A tensor of losses between model predictions and targets.
-        """
-        loss, preds, targets = self.model_step(batch)
-
-        # update and log metrics
-        self.train_loss(loss)
-        self.train_acc(preds, targets)
-        self.log("train/loss", self.train_loss, on_step=False, on_epoch=True, prog_bar=True)
-        self.log("train/acc", self.train_acc, on_step=False, on_epoch=True, prog_bar=True)
-
-        # return loss or backpropagation will fail
-        return loss
-
-    def on_train_epoch_end(self) -> None:
-        "Lightning hook that is called when a training epoch ends."
-        pass
-
-    def validation_step(self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> None:
-        """Perform a single validation step on a batch of data from the validation set.
-
-        :param batch: A batch of data (a tuple) containing the input tensor of images and target
-            labels.
-        :param batch_idx: The index of the current batch.
-        """
-        loss, preds, targets = self.model_step(batch)
-
-        # update and log metrics
-        self.val_loss(loss)
-        self.val_acc(preds, targets)
-        self.log("val/loss", self.val_loss, on_step=False, on_epoch=True, prog_bar=True)
-        self.log("val/acc", self.val_acc, on_step=False, on_epoch=True, prog_bar=True)
-
-    def on_validation_epoch_end(self) -> None:
-        "Lightning hook that is called when a validation epoch ends."
-        acc = self.val_acc.compute()  # get current val acc
-        self.val_acc_best(acc)  # update best so far val acc
-        # log `val_acc_best` as a value through `.compute()` method, instead of as a metric object
-        # otherwise metric would be reset by lightning after each epoch
-        self.log("val/acc_best", self.val_acc_best.compute(), sync_dist=True, prog_bar=True)
-
-    def test_step(self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> None:
-        """Perform a single test step on a batch of data from the test set.
-
-        :param batch: A batch of data (a tuple) containing the input tensor of images and target
-            labels.
-        :param batch_idx: The index of the current batch.
-        """
-        loss, preds, targets = self.model_step(batch)
-
-        # update and log metrics
-        self.test_loss(loss)
-        self.test_acc(preds, targets)
-        self.log("test/loss", self.test_loss, on_step=False, on_epoch=True, prog_bar=True)
-        self.log("test/acc", self.test_acc, on_step=False, on_epoch=True, prog_bar=True)
-
-    def on_test_epoch_end(self) -> None:
-        """Lightning hook that is called when a test epoch ends."""
-        pass
-
-    def setup(self, stage: str) -> None:
-        """Lightning hook that is called at the beginning of fit (train + validate), validate,
-        test, or predict.
-
-        This is a good hook when you need to build models dynamically or adjust something about
-        them. This hook is called on every process when using DDP.
-
-        :param stage: Either `"fit"`, `"validate"`, `"test"`, or `"predict"`.
-        """
-        if self.hparams.compile and stage == "fit":
-            self.net = torch.compile(self.net)
-
-    def configure_optimizers(self) -> Dict[str, Any]:
-        """Choose what optimizers and learning-rate schedulers to use in your optimization.
-        Normally you'd need one. But in the case of GANs or similar you might have multiple.
-
-        Examples:
-            https://lightning.ai/docs/pytorch/latest/common/lightning_module.html#configure-optimizers
-
-        :return: A dict containing the configured optimizers and learning-rate schedulers to be used for training.
-        """
-        optimizer = self.hparams.optimizer(params=self.trainer.model.parameters())
-        if self.hparams.scheduler is not None:
-            scheduler = self.hparams.scheduler(optimizer=optimizer)
-            return {
-                "optimizer": optimizer,
-                "lr_scheduler": {
-                    "scheduler": scheduler,
-                    "monitor": "val/loss",
-                    "interval": "epoch",
-                    "frequency": 1,
-                },
-            }
-        return {"optimizer": optimizer}
-
-
-if __name__ == "__main__":
-    _ = MNISTLitModule(None, None, None, None)

src/train.py

@@ -38,6 +38,8 @@ from src.utils import (
 
 log = RankedLogger(__name__, rank_zero_only=True)
 
+torch.set_float32_matmul_precision("medium")
+
 
 @task_wrapper
 def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
@@ -67,7 +69,9 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     logger: List[Logger] = instantiate_loggers(cfg.get("logger"))
 
     log.info(f"Instantiating trainer <{cfg.trainer._target_}>")
-    trainer: Trainer = hydra.utils.instantiate(cfg.trainer, callbacks=callbacks, logger=logger)
+    trainer: Trainer = hydra.utils.instantiate(
+        cfg.trainer, callbacks=callbacks, logger=logger
+    )
 
     object_dict = {
         "cfg": cfg,

test_bert.ipynb

@@ -0,0 +1,142 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/qninh/miniconda3/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "from src.data.mixed_datamodule import MixedDataModule"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "datamodule = MixedDataModule(dataset_path=\"./data/mixed\", batch_size=32, num_workers=4, bert_model=\"bert-base-uncased\", tool_capacity=16)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "datamodule.setup(stage=\"fit\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_dataloader = datamodule.train_dataloader()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'instruction': torch.Size([32, 128]),\n",
+       " 'instruction_mask': torch.Size([32, 128]),\n",
+       " 'tool_desc_emb': torch.Size([32, 128]),\n",
+       " 'tool_desc_mask': torch.Size([32, 128])}"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# first sample\n",
+    "batch = next(iter(train_dataloader))\n",
+    "{\n",
+    "    key: value.shape\n",
+    "    for key, value in batch.items()\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "val_dataloader = datamodule.val_dataloader()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'instruction': torch.Size([32, 128]),\n",
+       " 'instruction_mask': torch.Size([32, 128]),\n",
+       " 'tool_desc_emb': torch.Size([32, 16, 128]),\n",
+       " 'tool_desc_mask': torch.Size([32, 16, 128])}"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# first sample\n",
+    "batch = next(iter(val_dataloader))\n",
+    "{\n",
+    "    key: value.shape\n",
+    "    for key, value in batch.items()\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "swim",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}