Update config.json

config.json

@@ -1,6 +1,9 @@
 {
   "architectures": [
-    "JudgeModel"
+    "JudgeModel",
+    "JudgeCasualLMHead",
+    "JudgeWithQA",
+    "JudgeClassifier"
   ],
   "inference_mode": true,
   "batch_size": 32,
@@ -41,29 +44,38 @@
   "vocab_size": 50257,
   "warmup_steps": 1000,
   "weight_decay": 0.01,
-  "megatron_config": null,
-  "megatron_core": "megatron.core",
-  "modules_to_save": [
-    "classifier",
-    "score"
-  ],
-  "peft_type": "LORA",
-  "r": 16,
-  "rank_pattern": {},
-  "revision": null,
-  "target_modules": [
-    "c_proj",
-    "c_attn",
-    "c_fc",
-    "score"
-  ],
+
+  "task_heads": {
+    "lm_head": {
+      "type": "JudgeCasualLMHead",
+      "params": {
+        "vocab_size": 50257
+      }
+    },
+    "qa_head": {
+      "type": "JudgeWithQA",
+      "params": {
+        "num_labels": 2
+      }
+    },
+    "classifier_head": {
+      "type": "JudgeClassifier",
+      "params": {
+        "num_labels": 5
+      }
+    }
+  },
+
   "task_specific_params": {
     "text-generation": {
       "do_sample": true,
       "max_length": 100
+    },
+    "question-answering": {
+      "max_answer_length": 30
+    },
+    "sequence-classification": {
+      "eval_steps": 500
     }
-  },
-  "task_type": "SEQ_CLS",
-  "use_dora": false,
-  "use_rslora": false
+  }
 }