test: change config and dummy dataset for E2E CoCa test

Author: manasMauryax

config_files/config_example_coca.yaml

@@ -34,6 +34,7 @@ collate_fn:
     sample_keys:
       - images
       - ${settings.referencing_keys.sample_key}
+      - modality
     target_keys: []
     text_sample_key: ${settings.referencing_keys.sample_key}
     text_target_key: ${settings.referencing_keys.target_key}
@@ -50,6 +51,9 @@ train_dataset:
       - sample_key: input_ids
         sample_shape: [1024]
         sample_type: int
+      - sample_key: modality
+        sample_shape: [0]
+        sample_type: const
 
 val_dataset:
   component_key: dataset
@@ -63,6 +67,9 @@ val_dataset:
       - sample_key: input_ids
         sample_shape: [1024]
         sample_type: int
+      - sample_key: modality
+        sample_shape: [0]
+        sample_type: const
 
 train_dataloader:
   component_key: data_loader
@@ -174,13 +181,14 @@ model:
   variant_key: coca
   config:
     prediction_key: logits
-    vision_embd_prediction_key: vision_embeddings
+    modality_key: modality
+    modality_embd_prediction_key: modality_embeddings
     text_embd_prediction_key: text_embeddings
-    vision_cls_prediction_key: vision_cls
+    modality_cls_prediction_key: modality_cls
     text_cls_prediction_key: text_cls
-    vision_encoder_config:
+    modality_encoder_config:
       sample_key: images
-      prediction_key: vision_embeddings
+      prediction_key: modality_embeddings
       img_size: 224
       n_classes: Null # Disable vision transformer head
       n_layer: 12
@@ -212,6 +220,7 @@ model:
       epsilon: 1e-5
     n_pool_head: 8
     n_vision_queries: 256
+    n_audio_queries: Null
     bias_attn_pool: False
     epsilon_attn_pool: 1e-5
     weight_init:

src/modalities/dataloader/dataset.py

@@ -29,6 +29,7 @@ def _check_if_inbounds(self, idx: int):
 class DummySampleDataType(str, Enum):
     FLOAT = "float"
     INT = "int"
+    CONSTANT = "const"
 
 
 class DummySampleConfig(BaseModel):
@@ -53,6 +54,8 @@ def __init__(self, num_samples: int, sample_definition: Tuple[DummySampleConfig]
         self.num_samples = num_samples
         self.sample_definition = sample_definition
 
+        self.VISION = 1
+
     def __len__(self) -> int:
         return self.num_samples
 
@@ -66,6 +69,8 @@ def _create_random_sample(self):
                 data = np.random.randn(*s.sample_shape)
             elif s.sample_type == DummySampleDataType.INT:
                 data = np.random.randint(low=0, high=512, size=s.sample_shape)
+            elif s.sample_type == DummySampleDataType.CONSTANT:
+                data = self.VISION
             else:
                 raise NotImplementedError(f"DummyDataset does not support type { s.sample_type}")
             sample[s.sample_key] = data