fix new hf dataset

piercus · Feb 17, 2024 · a2e5639 · a2e5639
1 parent 809d85f
commit a2e5639
Show file tree

Hide file tree

Showing 3 changed files with 1 addition and 5 deletions.
diff --git a/configs/remote/scheduled/finetune-color-palette-mlp-unweighted.toml b/configs/remote/scheduled/finetune-color-palette-mlp-unweighted.toml
@@ -60,7 +60,6 @@ use_gyro_dropout = false
 hf_repo = "refiners/unsplash-25k-cogvlm-captions"
 revision = "main"
 resize_image_max_size = 512
-caption_key = "ai_description"
 split = "train[200:]"
 
 [checkpointing]
@@ -71,7 +70,6 @@ use_wandb = true
 hf_repo = "refiners/unsplash-25k-cogvlm-captions"
 revision = "main"
 resize_image_max_size = 512
-caption_key = "ai_description"
 split = "train[:20]"
 
 [evaluation]

diff --git a/configs/remote/scheduled/finetune-color-palette-mlp-weighted.toml b/configs/remote/scheduled/finetune-color-palette-mlp-weighted.toml
@@ -60,7 +60,6 @@ use_gyro_dropout = false
 hf_repo = "refiners/unsplash-25k-cogvlm-captions"
 revision = "main"
 resize_image_max_size = 512
-caption_key = "ai_description"
 split = "train[200:]"
 
 [checkpointing]
@@ -71,7 +70,6 @@ use_wandb = true
 hf_repo = "refiners/unsplash-25k-cogvlm-captions"
 revision = "main"
 resize_image_max_size = 512
-caption_key = "ai_description"
 split = "train[:20]"
 
 [evaluation]

diff --git a/src/refiners/training_utils/huggingface_datasets.py b/src/refiners/training_utils/huggingface_datasets.py
@@ -29,7 +29,7 @@ def load_hf_dataset(
 
     dataset = dataset.map(
         function=download_image,
-        input_columns=["url"],
+        input_columns=["photo_image_url"],
         fn_kwargs={
             "dl_manager": DownloadManager(),
         },