Update training

src/data/basic_dataset.py

@@ -76,6 +76,9 @@ def __init__(
         data = load_dataset(split)
 
         if split == "train":
+            if add_bev:
+                data = get_calibration(data)
+                data = generate_bevs(data, front_cam=front_cam)
             data = normalise_key_object_infos(data, resize_factor, use_grid)
 
         if split == "train" and add_augmented:

src/reasoning/reasoning_engine.py

@@ -37,7 +37,7 @@ def process_batch(self, batch_items: List[QueryItem]) -> List[QueryItem]:
                     image_path=item.image_path,
                     qa_id=f"{item.qa_id}_reasoning",
                     qa_type=item.qa_type,
-                    key_object_info=item.key_object_info,  # note not available in eval mode
+                    key_object_info=item.key_object_info,
                     system_prompt=item.system_prompt,
                 )
                 desc_item.formatted_message = desc_item.format_message(

src/train/train_qwen.py

@@ -41,7 +41,7 @@ class TrainingArguments(transformers.TrainingArguments):
     cache_dir: Optional[str] = field(default=None)
     optim: str = field(default="adamw_torch")
     model_max_length: int = field(
-        default=512,
+        default=1028,
         metadata={
             "help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
         },
@@ -242,7 +242,6 @@ def create_optimizer(self):
     return self.optimizer
 
 
-# TODO: Look into the deepspeed config
 def train(
     approach_name: str,
     resize_factor: float,
@@ -252,6 +251,9 @@ def train(
     use_augmented: bool = False,
     use_reasoning: bool = False,
     use_system_prompt: bool = False,
+    add_kois: bool = False,
+    add_bev: bool = False,
+    front_cam: bool = False,
     **kwargs,
 ):
     name = approach_name + datetime.now().strftime("%H:%M:%S-%m-%d-%Y%")
@@ -310,14 +312,17 @@ def collator(batch: Any):
     dataset = DriveLMImageDataset(
         engine.training_message_formatter,
         split="train",
+        front_cam=front_cam,
+        add_kois=add_kois,
+        add_bev=add_bev,
         use_grid=use_grid,
         add_augmented=use_augmented,
         use_reasoning=use_reasoning,
         use_system_prompt=use_system_prompt,
         resize_factor=resize_factor,
     )
     if test_set_size is not None:
-        dataset = create_subset(dataset, int(test_set_size))
+        dataset = create_subset(dataset, int(test_set_size), equal_distribution=True)
     dataset = [item.formatted_message for item in dataset]
 
     engine.load_model(flash_attn=False)