add wandb to pre_processing

pre_processing.py

@@ -10,6 +10,7 @@ from transformers import BertTokenizer
 from constants import (RAW_DATA_DIR,
                        PROCESSED_DATA_DIR,
                        METADATA_FILEPATH,
+                       PROJECT_NAME,
                        BERT_BASE,
                        MAX_SEQUENCE_LENGHT,
                        FilePath,
@@ -18,6 +19,8 @@ from constants import (RAW_DATA_DIR,
 # Allow for unlimited image size, some documents are pretty big...
 Image.MAX_IMAGE_PIXELS = None
 
+run = wandb.init(project=PROJECT_NAME, name='pre-processing')
+
 
 def make_page_filepaths(basename, label, page_index) -> Tuple[str, str]:
     out_dirname = path.join(PROCESSED_DATA_DIR, label)
@@ -108,4 +111,21 @@ def process_training_data() -> pd.DataFrame:
     return pages_metadata_df
 
 
-process_training_data()
+def main():
+    metadata_df = process_training_data()
+
+    raw_dataset_artifact = wandb.Artifact("raw-dataset", type="dataset")
+    raw_dataset_artifact.add_dir(RAW_DATA_DIR)
+    run.log_artifact(raw_dataset_artifacth)
+
+    processed_dataset_artifact = wandb.Artifact("processed-dataset", type="dataset")
+    processed_dataset_artifact.add_dir(PROCESSED_DATA_DIR)
+    run.log_artifact(processed_dataset_artifact)
+
+    dataset_metadata_artifact = wandb.Artifact("dataset-metadata", type="dataset")
+    dataset_metadata_table = wandb.Table(dataframe=metadata_df)
+    dataset_metadata_artifact.add(dataset_metadata_table, name='metadata-table')
+    run.log_artifact(dataset_metadata_artifact)
+
+if __name__ == '__main__':
+    main()