dummy wandb to training code

training.py

@@ -90,36 +90,28 @@ def prepare_dataset(
         .prefetch(tf.data.experimental.AUTOTUNE)
 
 
-metadata_df: DataFrame[PageMetadata] = pd.read_csv(METADATA_FILEPATH)
-metadata_df = metadata_df.sample(n=50, random_state=42)
-
-median_height = int(metadata_df['height'].median())
-median_width = int(metadata_df['width'].median())
-img_size: ImageSize = (median_height, median_width)
-img_input_shape: ImageInputShape = img_size + (3,)
-
-label_names: List[str] = sorted(
-    [d.name for d in PROCESSED_DATA_DIR.iterdir() if d.is_dir()]
-)
-num_classes = len(label_names)
+def prepare_data(
+        df: DataFrame[PageMetadata]
+) -> Tuple[tf.data.Dataset, tf.data.Dataset, tf.data.Dataset]:
+    print('Splitting the DataFrame into training, validation and test')
+    train_df, val_df, test_df = stratified_split(
+        df,
+        train_frac=0.7,
+        val_frac=0.15,
+        test_frac=0.15,
+    )
 
-print('Splitting the DataFrame into training, validation and test')
-train_df, val_df, test_df = stratified_split(
-    metadata_df,
-    train_frac=0.7,
-    val_frac=0.15,
-    test_frac=0.15,
-)
+    print('Batching and shuffling the datasets')
+    train_ds = dataset_from_dataframe(train_df)
+    train_ds = prepare_dataset(train_ds, img_size, batch_size=BATCH_SIZE)
 
-print('Batching and shuffling the datasets')
-train_ds = dataset_from_dataframe(train_df)
-train_ds = prepare_dataset(train_ds, img_size, batch_size=BATCH_SIZE)
+    val_ds = dataset_from_dataframe(val_df)
+    val_ds = prepare_dataset(val_ds, img_size, batch_size=BATCH_SIZE)
 
-val_ds = dataset_from_dataframe(val_df)
-val_ds = prepare_dataset(val_ds, img_size, batch_size=BATCH_SIZE)
+    test_ds = dataset_from_dataframe(test_df)
+    test_ds = prepare_dataset(test_ds, img_size, batch_size=BATCH_SIZE)
 
-test_ds = dataset_from_dataframe(test_df)
-test_ds = prepare_dataset(test_ds, img_size, batch_size=BATCH_SIZE)
+    return train_ds, val_ds, test_ds
 
 
 def build_image_model(input_shape: ImageInputShape) -> keras.Model:
@@ -199,16 +191,41 @@ def build_multimodal_model(
     return multimodal_model
 
 
-multimodal_model = build_multimodal_model(num_classes, img_input_shape)
-multimodal_model.summary()
-multimodal_model.compile(
-    optimizer='adam',
-    loss='sparse_categorical_crossentropy',
-    metrics=['accuracy']
-)
-multimodal_model.fit(
-    train_ds,
-    epochs=EPOCHS,
-    batch_size=BATCH_SIZE,
-    validation_data=val_ds,
-)
+def train():
+    metadata_df: DataFrame[PageMetadata] = pd.read_csv(METADATA_FILEPATH)
+
+    median_height = int(metadata_df['height'].median())
+    median_width = int(metadata_df['width'].median())
+
+    img_size: ImageSize = (median_height, median_width)
+    img_input_shape: ImageInputShape = img_size + (3,)
+
+    label_names: List[str] = sorted(
+        [d.name for d in PROCESSED_DATA_DIR.iterdir() if d.is_dir()]
+    )
+    num_classes = len(label_names)
+
+    train_ds, val_ds, test_ds = prepare_data(metadata_df)
+
+    multimodal_model = build_multimodal_model(num_classes, img_input_shape)
+    multimodal_model.summary()
+    multimodal_model.compile(
+        optimizer='adam',
+        loss='sparse_categorical_crossentropy',
+        metrics=['accuracy']
+    )
+    multimodal_model.fit(
+        train_ds,
+        epochs=EPOCHS,
+        batch_size=BATCH_SIZE,
+        validation_data=val_ds,
+    )
+
+
+def evaluate():
+    return
+
+
+if __name__ = '__main__':
+    train()
+    evaluate()