fix: add sample dataset and missing enum

ruflab · Oct 22, 2024 · df4bbce · df4bbce
1 parent bdbaf37
commit df4bbce
Show file tree

Hide file tree

Showing 41 changed files with 24 additions and 3 deletions.
diff --git a/sample_dataset/domain_splits/test_attr:1.0_attr,v:0.5_v:1.0_seed:0_domain_split.npy b/sample_dataset/domain_splits/test_attr:1.0_attr,v:0.5_v:1.0_seed:0_domain_split.npy
diff --git a/sample_dataset/domain_splits/test_t:1.0_t,v:0.5_v:1.0_seed:0_domain_split.npy b/sample_dataset/domain_splits/test_t:1.0_t,v:0.5_v:1.0_seed:0_domain_split.npy
diff --git a/sample_dataset/domain_splits/train_attr:1.0_attr,v:0.5_v:1.0_seed:0_domain_split.npy b/sample_dataset/domain_splits/train_attr:1.0_attr,v:0.5_v:1.0_seed:0_domain_split.npy
diff --git a/sample_dataset/domain_splits/train_t:1.0_t,v:0.5_v:1.0_seed:0_domain_split.npy b/sample_dataset/domain_splits/train_t:1.0_t,v:0.5_v:1.0_seed:0_domain_split.npy
diff --git a/sample_dataset/domain_splits/val_attr:1.0_attr,v:0.5_v:1.0_seed:0_domain_split.npy b/sample_dataset/domain_splits/val_attr:1.0_attr,v:0.5_v:1.0_seed:0_domain_split.npy
diff --git a/sample_dataset/domain_splits/val_t:1.0_t,v:0.5_v:1.0_seed:0_domain_split.npy b/sample_dataset/domain_splits/val_t:1.0_t,v:0.5_v:1.0_seed:0_domain_split.npy
diff --git a/sample_dataset/domain_splits_v2/test_attr:1.0_attr,v:0.5_v:1.0_seed:0_ms:2_domain_split.npy b/sample_dataset/domain_splits_v2/test_attr:1.0_attr,v:0.5_v:1.0_seed:0_ms:2_domain_split.npy
diff --git a/sample_dataset/domain_splits_v2/test_t:1.0_t,v:0.5_v:1.0_seed:0_ms:2_domain_split.npy b/sample_dataset/domain_splits_v2/test_t:1.0_t,v:0.5_v:1.0_seed:0_ms:2_domain_split.npy
diff --git a/sample_dataset/domain_splits_v2/train_attr:1.0_attr,v:0.5_v:1.0_seed:0_ms:4_domain_split.npy b/sample_dataset/domain_splits_v2/train_attr:1.0_attr,v:0.5_v:1.0_seed:0_ms:4_domain_split.npy
diff --git a/sample_dataset/domain_splits_v2/train_t:1.0_t,v:0.5_v:1.0_seed:0_ms:4_domain_split.npy b/sample_dataset/domain_splits_v2/train_t:1.0_t,v:0.5_v:1.0_seed:0_ms:4_domain_split.npy
diff --git a/sample_dataset/domain_splits_v2/val_attr:1.0_attr,v:0.5_v:1.0_seed:0_ms:2_domain_split.npy b/sample_dataset/domain_splits_v2/val_attr:1.0_attr,v:0.5_v:1.0_seed:0_ms:2_domain_split.npy
diff --git a/sample_dataset/domain_splits_v2/val_t:1.0_t,v:0.5_v:1.0_seed:0_ms:2_domain_split.npy b/sample_dataset/domain_splits_v2/val_t:1.0_t,v:0.5_v:1.0_seed:0_ms:2_domain_split.npy
diff --git a/sample_dataset/latent_mean.npy b/sample_dataset/latent_mean.npy
diff --git a/sample_dataset/latent_std.npy b/sample_dataset/latent_std.npy
diff --git a/sample_dataset/test/0.png b/sample_dataset/test/0.png
diff --git a/sample_dataset/test/1.png b/sample_dataset/test/1.png
diff --git a/sample_dataset/test_caption_choices.npy b/sample_dataset/test_caption_choices.npy
diff --git a/sample_dataset/test_captions.npy b/sample_dataset/test_captions.npy
diff --git a/sample_dataset/test_labels.npy b/sample_dataset/test_labels.npy
diff --git a/sample_dataset/test_latent.npy b/sample_dataset/test_latent.npy
diff --git a/sample_dataset/test_odd_one_out_labels.npy b/sample_dataset/test_odd_one_out_labels.npy
diff --git a/sample_dataset/test_unpaired.npy b/sample_dataset/test_unpaired.npy
diff --git a/sample_dataset/train/0.png b/sample_dataset/train/0.png
diff --git a/sample_dataset/train/1.png b/sample_dataset/train/1.png
diff --git a/sample_dataset/train/2.png b/sample_dataset/train/2.png
diff --git a/sample_dataset/train/3.png b/sample_dataset/train/3.png
diff --git a/sample_dataset/train_caption_choices.npy b/sample_dataset/train_caption_choices.npy
diff --git a/sample_dataset/train_captions.npy b/sample_dataset/train_captions.npy
diff --git a/sample_dataset/train_labels.npy b/sample_dataset/train_labels.npy
diff --git a/sample_dataset/train_latent.npy b/sample_dataset/train_latent.npy
diff --git a/sample_dataset/train_odd_one_out_labels.npy b/sample_dataset/train_odd_one_out_labels.npy
diff --git a/sample_dataset/train_unpaired.npy b/sample_dataset/train_unpaired.npy
diff --git a/sample_dataset/val/0.png b/sample_dataset/val/0.png
diff --git a/sample_dataset/val/1.png b/sample_dataset/val/1.png
diff --git a/sample_dataset/val_caption_choices.npy b/sample_dataset/val_caption_choices.npy
diff --git a/sample_dataset/val_captions.npy b/sample_dataset/val_captions.npy
diff --git a/sample_dataset/val_labels.npy b/sample_dataset/val_labels.npy
diff --git a/sample_dataset/val_latent.npy b/sample_dataset/val_latent.npy
diff --git a/sample_dataset/val_odd_one_out_labels.npy b/sample_dataset/val_odd_one_out_labels.npy
diff --git a/sample_dataset/val_unpaired.npy b/sample_dataset/val_unpaired.npy
diff --git a/simple_shapes_dataset/dataset/domain.py b/simple_shapes_dataset/dataset/domain.py
@@ -1,4 +1,5 @@
 from collections.abc import Callable, Iterable
+from enum import Enum
 from pathlib import Path
 from typing import Any, NamedTuple, TypedDict
 
@@ -7,7 +8,26 @@
 from PIL import Image
 from shimmer import DataDomain, DomainDesc
 
-from simple_shapes_dataset.types import DomainType
+
+class DomainType(Enum):
+    v = DomainDesc("v", "v")
+    v_latents = DomainDesc("v", "v_latents")
+    attr = DomainDesc("attr", "attr")
+    t = DomainDesc("t", "t")
+    raw_text = DomainDesc("t", "raw_text")
+
+
+class DomainModelVariantType(Enum):
+    v = (DomainType.v, "default")
+    attr = (DomainType.attr, "default")
+    attr_legacy = (DomainType.attr, "legacy")
+    attr_unpaired = (DomainType.attr, "unpaired")
+    v_latents = (DomainType.v_latents, "default")
+    v_latents_unpaired = (DomainType.v_latents, "unpaired")
+
+    def __init__(self, kind: DomainType, model_variant: str) -> None:
+        self.kind = kind
+        self.model_variant = model_variant
 
 
 class SimpleShapesImages(DataDomain):
@@ -305,6 +325,7 @@ def get_default_domains(
 ) -> dict[DomainDesc, type[DataDomain]]:
     domain_classes = {}
     for domain in domains:
-        domain_desc = DomainType[domain].value if isinstance(domain, str) else domain
-        domain_classes[domain_desc] = DEFAULT_DOMAINS[domain_desc.kind]
+        if isinstance(domain, str):
+            domain = DomainType[domain].value
+        domain_classes[domain] = DEFAULT_DOMAINS[domain.kind]
     return domain_classes