Update

app_lib/config.py

@@ -0,0 +1,159 @@
+# SOURCE: https://github.com/Sulam-Group/IBYDMT/blob/main/ibydmt/utils/config.py
+
+import os
+from dataclasses import dataclass
+from enum import Enum
+from itertools import product
+from typing import Any, Iterable, Mapping, Optional, Union
+
+import torch
+from ml_collections import ConfigDict
+from numpy import ndarray
+
+Array = Union[ndarray, torch.Tensor]
+
+
+class TestType(Enum):
+    GLOBAL = "global"
+    GLOBAL_COND = "global_cond"
+    LOCAL_COND = "local_cond"
+
+
+class ConceptType(Enum):
+    DATASET = "dataset"
+    CLASS = "class"
+    IMAGE = "image"
+
+
+@dataclass
+class Constants:
+    WORKDIR = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
+    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class DataConfig(ConfigDict):
+    def __init__(self, config_dict: Optional[Mapping[str, Any]] = None):
+        super().__init__()
+        if config_dict is None:
+            config_dict = {}
+
+        self.dataset: str = config_dict.get("dataset", None)
+        self.backbone: str = config_dict.get("backbone", None)
+        self.bottleneck: str = config_dict.get("bottleneck", None)
+        self.classifier: str = config_dict.get("classifier", None)
+        self.sampler: str = config_dict.get("sampler", None)
+        self.num_concepts: int = config_dict.get("num_concepts", None)
+
+
+class SpliceConfig(ConfigDict):
+    def __init__(self, config_dict: Optional[Mapping[str, Any]] = None):
+        super().__init__()
+        if config_dict is None:
+            config_dict = {}
+
+        self.vocab: str = config_dict.get("vocab", None)
+        self.vocab_size: int = config_dict.get("vocab_size", None)
+        self.l1_penalty: float = config_dict.get("l1_penalty", None)
+
+
+class PCBMConfig(ConfigDict):
+    def __init__(self, config_dict: Optional[Mapping[str, Any]] = None):
+        super().__init__()
+        if config_dict is None:
+            config_dict = {}
+
+        self.alpha: float = config_dict.get("alpha", None)
+        self.l1_ratio: float = config_dict.get("l1_ratio", None)
+
+
+class cKDEConfig(ConfigDict):
+    def __init__(self, config_dict: Optional[Mapping[str, Any]] = None):
+        super().__init__()
+        if config_dict is None:
+            config_dict = {}
+
+        self.metric: str = config_dict.get("metric", None)
+        self.scale_method: str = config_dict.get("scale_method", None)
+        self.scale: float = config_dict.get("scale", None)
+
+
+class TestingConfig(ConfigDict):
+    def __init__(self, config_dict: Optional[Mapping[str, Any]] = None):
+        super().__init__()
+        if config_dict is None:
+            config_dict = {}
+
+        self.significance_level: float = config_dict.get("significance_level", None)
+        self.wealth: str = config_dict.get("wealth", None)
+        self.bet: str = config_dict.get("bet", None)
+        self.kernel: str = config_dict.get("kernel", None)
+        self.kernel_scale_method: str = config_dict.get("kernel_scale_method", None)
+        self.kernel_scale: float = config_dict.get("kernel_scale", None)
+        self.tau_max: int = config_dict.get("tau_max", None)
+        self.images_per_class: int = config_dict.get("images_per_class", None)
+        self.r: int = config_dict.get("r", None)
+        self.cardinality: Iterable[int] = config_dict.get("cardinality", None)
+
+
+class Config(ConfigDict):
+    def __init__(self, config_dict: Optional[Mapping[str, Any]] = None):
+        super().__init__()
+        if config_dict is None:
+            config_dict = {}
+
+        self.name: str = config_dict.get("name", None)
+        self.data = DataConfig(config_dict.get("data", None))
+        self.splice = SpliceConfig(config_dict.get("splice", None))
+        self.pcbm = PCBMConfig(config_dict.get("pcbm", None))
+        self.ckde = cKDEConfig(config_dict.get("ckde", None))
+        self.testing = TestingConfig(config_dict.get("testing", None))
+
+    def backbone_name(self):
+        backbone = self.data.backbone.strip().lower()
+        return backbone.replace("/", "_").replace(":", "_")
+
+    def sweep(self, keys: Iterable[str]):
+        def _get(dict, key):
+            keys = key.split(".")
+            if len(keys) == 1:
+                return dict[keys[0]]
+            else:
+                return _get(dict[keys[0]], ".".join(keys[1:]))
+
+        def _set(dict, key, value):
+            keys = key.split(".")
+            if len(keys) == 1:
+                dict[keys[0]] = value
+            else:
+                _set(dict[keys[0]], ".".join(keys[1:]), value)
+
+        to_iterable = lambda v: v if isinstance(v, list) else [v]
+
+        config_dict = self.to_dict()
+        sweep_values = [_get(config_dict, key) for key in keys]
+        sweep = list(product(*map(to_iterable, sweep_values)))
+
+        configs: Iterable[Config] = []
+        for _sweep in sweep:
+            _config_dict = config_dict.copy()
+            for key, value in zip(keys, _sweep):
+                _set(_config_dict, key, value)
+
+            configs.append(Config(_config_dict))
+        return configs
+
+
+def register_config(name: str):
+    def register(cls: Config):
+        if name in configs:
+            raise ValueError(f"Config {name} is already registered")
+        configs[name] = cls
+
+    return register
+
+
+def get_config(name: str) -> Config:
+    return configs[name]()
+
+
+configs: Mapping[str, Config] = {}

app_lib/defaults.py

@@ -1,14 +1,19 @@
-DATASET_NAME = "imagenette"
-MODEL_NAME = "open_clip:ViT-B-32"
+from dataclasses import dataclass
 
-SIGNIFICANCE_LEVEL_VALUE = 0.05
-SIGNIFICANCE_LEVEL_STEP = 0.01
 
-TAU_MAX_VALUE = 200
-TAU_MAX_STEP = 50
+@dataclass
+class Defaults:
+    DATASET_NAME = "imagenette"
+    MODEL_NAME = "open_clip:ViT-B-32"
 
-R_VALUE = 20
-R_STEP = 5
+    SIGNIFICANCE_LEVEL_VALUE = 0.05
+    SIGNIFICANCE_LEVEL_STEP = 0.01
 
-CARDINALITY_VALUE = 1
-CARDINALITY_STEP = 1
+    TAU_MAX_VALUE = 200
+    TAU_MAX_STEP = 50
+
+    R_VALUE = 20
+    R_STEP = 5
+
+    CARDINALITY_VALUE = 1
+    CARDINALITY_STEP = 1

app_lib/multimodal.py

@@ -0,0 +1,186 @@
+# SOURCE: https://github.com/Sulam-Group/IBYDMT/blob/main/ibydmt/multimodal.py
+
+from abc import abstractmethod
+from typing import Mapping, Optional
+
+import clip
+import open_clip
+from transformers import (
+    AlignModel,
+    AlignProcessor,
+    BlipForImageTextRetrieval,
+    BlipProcessor,
+    FlavaModel,
+    FlavaProcessor,
+)
+
+from app_lib.config import Config
+from app_lib.config import Constants as c
+
+
+class VisionLanguageModel:
+    def __init__(self, backbone: Optional[str] = None, device=c.DEVICE):
+        pass
+
+    @abstractmethod
+    def encode_text(self, text):
+        pass
+
+    @abstractmethod
+    def encode_image(self, image):
+        pass
+
+
+models: Mapping[str, VisionLanguageModel] = {}
+
+
+def register_model(name):
+    def register(cls: VisionLanguageModel):
+        if name in models:
+            raise ValueError(f"Model {name} is already registered")
+        models[name] = cls
+
+    return register
+
+
+def get_model_name_and_backbone(config: Config):
+    backbone = config.data.backbone.split(":")
+    if len(backbone) == 1:
+        backbone.append(None)
+    return backbone
+
+
+def get_model(config: Config, device=c.DEVICE) -> VisionLanguageModel:
+    model_name, backbone = get_model_name_and_backbone(config)
+    return models[model_name](backbone, device=device)
+
+
+def get_text_encoder(config: Config, device=c.DEVICE):
+    model = get_model(config, device=device)
+    return model.encode_text
+
+
+def get_image_encoder(config: Config, device=c.DEVICE):
+    model = get_model(config, device=device)
+    return model.encode_image
+
+
+@register_model(name="clip")
+class CLIPModel(VisionLanguageModel):
+    def __init__(self, backbone: str, device=c.DEVICE):
+        self.model, self.preprocess = clip.load(backbone, device=device)
+        self.tokenize = clip.tokenize
+
+        self.device = device
+
+    def encode_text(self, text):
+        text = self.tokenize(text).to(self.device)
+        return self.model.encode_text(text)
+
+    def encode_image(self, image):
+        image = self.preprocess(image).unsqueeze(0).to(self.device)
+        return self.model.encode_image(image)
+
+
+@register_model(name="open_clip")
+class OpenClipModel(VisionLanguageModel):
+    OPENCLIP_WEIGHTS = {
+        "ViT-B-32": "laion2b_s34b_b79k",
+        "ViT-L-14": "laion2b_s32b_b82k",
+    }
+
+    def __init__(self, backbone: str, device=c.DEVICE):
+        self.model, _, self.preprocess = open_clip.create_model_and_transforms(
+            backbone, pretrained=self.OPENCLIP_WEIGHTS[backbone], device=device
+        )
+        self.tokenize = open_clip.get_tokenizer(backbone)
+
+        self.device = device
+
+    def encode_text(self, text):
+        text = self.tokenize(text).to(self.device)
+        return self.model.encode_text(text)
+
+    def encode_image(self, image):
+        image = self.preprocess(image).unsqueeze(0).to(self.device)
+        return self.model.encode_image(image)
+
+
+@register_model(name="flava")
+class FLAVAModel(VisionLanguageModel):
+    HF_MODEL = "facebook/flava-full"
+
+    def __init__(self, backbone: Optional[str] = None, device=c.DEVICE):
+        if backbone is None:
+            backbone = self.HF_MODEL
+
+        self.model = FlavaModel.from_pretrained(backbone).to(device)
+        self.processor = FlavaProcessor.from_pretrained(backbone)
+
+        self.device = device
+
+    def encode_text(self, text):
+        text_inputs = self.processor(
+            text=text, return_tensors="pt", padding="max_length", max_length=77
+        )
+        text_inputs = {k: v.to(self.device) for k, v in text_inputs.items()}
+        return self.model.get_text_features(**text_inputs)[:, 0, :]
+
+    def encode_image(self, image):
+        image_inputs = self.processor(images=image, return_tensors="pt")
+        image_inputs = {k: v.to(self.device) for k, v in image_inputs.items()}
+        return self.model.get_image_features(**image_inputs)[:, 0, :]
+
+
+@register_model(name="align")
+class ALIGNModel(VisionLanguageModel):
+    HF_MODEL = "kakaobrain/align-base"
+
+    def __init__(self, backbone: Optional[str] = None, device=c.DEVICE):
+        if backbone is None:
+            backbone = self.HF_MODEL
+
+        self.model = AlignModel.from_pretrained(backbone).to(device)
+        self.processor = AlignProcessor.from_pretrained(backbone)
+
+        self.device = device
+
+    def encode_text(self, text):
+        text_inputs = self.processor(
+            text=text, return_tensors="pt", padding="max_length", max_length=77
+        )
+        text_inputs = {k: v.to(self.device) for k, v in text_inputs.items()}
+        return self.model.get_text_features(**text_inputs)
+
+    def encode_image(self, image):
+        image_inputs = self.processor(images=image, return_tensors="pt")
+        image_inputs = {k: v.to(self.device) for k, v in image_inputs.items()}
+        return self.model.get_image_features(**image_inputs)
+
+
+@register_model(name="blip")
+class BLIPModel(VisionLanguageModel):
+    HF_MODEL = "Salesforce/blip-itm-base-coco"
+
+    def __init__(self, backbone: Optional[str] = None, device=c.DEVICE):
+        if backbone is None:
+            backbone = self.HF_MODEL
+
+        self.model = BlipForImageTextRetrieval.from_pretrained(backbone).to(device)
+        self.processor = BlipProcessor.from_pretrained(backbone)
+
+        self.device = device
+
+    def encode_text(self, text):
+        text_inputs = self.processor(
+            text=text, return_tensors="pt", padding="max_length", max_length=77
+        )
+        text_inputs = {k: v.to(self.device) for k, v in text_inputs.items()}
+        question_embeds = self.model.text_encoder(**text_inputs)[0]
+        return self.model.text_proj(question_embeds[:, 0, :])
+
+    def encode_image(self, image):
+        image_inputs = self.processor(images=image, return_tensors="pt")
+        image_inputs = {k: v.to(self.device) for k, v in image_inputs.items()}
+        image_embeds = self.model.vision_model(**image_inputs)[0]
+        return self.model.vision_proj(image_embeds[:, 0, :])

app_lib/test.py

@@ -1,65 +1,33 @@
 import os
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
-import clip
-import h5py
 import ml_collections
 import numpy as np
-import open_clip
+import pandas as pd
 import streamlit as st
 import torch
 from huggingface_hub import hf_hub_download
 
+import app_lib.multimodal as multimodal
 from app_lib.ckde import cKDE
-from app_lib.utils import SUPPORTED_MODELS
+from app_lib.config import Config
+from app_lib.config import Constants as c
 from ibydmt.test import xSKIT
 
 rng = np.random.default_rng()
 
 
-def _get_open_clip_model(model_name, device):
-    backbone = model_name.split(":")[-1]
-
-    model, _, preprocess = open_clip.create_model_and_transforms(
-        SUPPORTED_MODELS[model_name], device=device
-    )
-    model.eval()
-    tokenizer = open_clip.get_tokenizer(backbone)
-    return model, preprocess, tokenizer
-
-
-def _get_clip_model(model_name, device):
-    backbone = model_name.split(":")[-1]
-    model, preprocess = clip.load(backbone, device=device)
-    tokenizer = clip.tokenize
-    return model, preprocess, tokenizer
-
-
-def _load_model(model_name, device):
-    if "open_clip" in model_name:
-        model, preprocess, tokenizer = _get_open_clip_model(model_name, device)
-    elif "clip" in model_name:
-        model, preprocess, tokenizer = _get_clip_model(model_name, device)
-    return model, preprocess, tokenizer
-
-
 @torch.no_grad()
 @torch.cuda.amp.autocast()
-def _encode_concepts(tokenizer, model, concepts, device):
-    concepts_text = tokenizer(concepts).to(device)
-
-    concept_features = model.encode_text(concepts_text)
+def _encode_concepts(model, concepts):
+    concept_features = model.encode_text(concepts)
     concept_features /= torch.linalg.norm(concept_features, dim=-1, keepdim=True)
     return concept_features.cpu().numpy()
 
 
 @torch.no_grad()
 @torch.cuda.amp.autocast()
-def _encode_image(model, preprocess, image, device):
-    image = preprocess(image)
-    image = image.unsqueeze(0)
-    image = image.to(device)
-
+def _encode_image(model, image):
     image_features = model.encode_image(image)
     image_features /= image_features.norm(dim=-1, keepdim=True)
     return image_features.cpu().numpy()
@@ -67,24 +35,24 @@ def _encode_image(model, preprocess, image, device):
 
 @torch.no_grad()
 @torch.cuda.amp.autocast()
-def _encode_class_name(tokenizer, model, class_name, device):
-    class_text = tokenizer([f"A photo of a {class_name}"]).to(device)
-
+def _encode_class_name(model, class_name):
+    class_text = [f"A photo of a {class_name}"]
     class_features = model.encode_text(class_text)
     class_features /= torch.linalg.norm(class_features, dim=-1, keepdim=True)
     return class_features.cpu().numpy()
 
 
-def _load_dataset(dataset_name, model_name):
+def _load_embedding(config):
     dataset_path = hf_hub_download(
         repo_id="jacopoteneggi/IBYDMT",
-        filename=f"{dataset_name}_{model_name}_train.h5",
+        filename=(
+            f"{config.data.dataset.lower()}_train_{config.backbone_name()}.parquet"
+        ),
         repo_type="dataset",
     )
 
-    with h5py.File(dataset_path, "r") as dataset:
-        embedding = dataset["embedding"][:]
-    return embedding
+    dataset = pd.read_parquet(dataset_path)
+    return np.array(dataset["embedding"].values.tolist())
 
 
 def _sample_random_subset(concept_idx, concepts, cardinality):
@@ -162,26 +130,30 @@ def test(
     cardinality,
     dataset_name,
     model_name,
-    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+    device=c.DEVICE,
     with_streamlit=True,
 ):
+    config = Config()
+    config.data.dataset = dataset_name
+    config.data.backbone = model_name
+
     if with_streamlit:
         with st.spinner("Loading model"):
-            model, preprocess, tokenizer = _load_model(model_name, device)
+            model = multimodal.get_model(config, device=device)
     else:
-        model, preprocess, tokenizer = _load_model(model_name, device)
+        model = multimodal.get_model(config, device=device)
 
     if with_streamlit:
         with st.spinner("Encoding concepts"):
-            cbm = _encode_concepts(tokenizer, model, concepts, device)
+            cbm = _encode_concepts(model, concepts)
     else:
-        cbm = _encode_concepts(tokenizer, model, concepts, device)
+        cbm = _encode_concepts(model, concepts)
 
     if with_streamlit:
         with st.spinner("Encoding image"):
-            h = _encode_image(model, preprocess, image, device)
+            h = _encode_image(model, image)
     else:
-        h = _encode_image(model, preprocess, image, device)
+        h = _encode_image(model, image)
     z = h @ cbm.T
     z = z.squeeze()
 
@@ -201,11 +173,11 @@ def test(
             ),
         )
 
-    embedding = _load_dataset(dataset_name, model_name)
+    embedding = _load_embedding(config)
     semantics = embedding @ cbm.T
     sampler = cKDE(embedding, semantics)
 
-    classifier = _encode_class_name(tokenizer, model, class_name, device)
+    classifier = _encode_class_name(model, class_name)
 
     with ThreadPoolExecutor() as executor:
         futures = [

app_lib/user_input.py

@@ -5,7 +5,7 @@ import streamlit as st
 from PIL import Image
 from streamlit_image_select import image_select
 
-import app_lib.defaults as defaults
+from app_lib.defaults import Defaults as d
 from app_lib.utils import SUPPORTED_DATASETS, SUPPORTED_MODELS
 
 IMAGE_DIR = os.path.join("assets", "images")
@@ -31,8 +31,8 @@ def _validate_concepts(concepts):
 
 
 def _get_significance_level():
-    default = defaults.SIGNIFICANCE_LEVEL_VALUE
-    step = defaults.SIGNIFICANCE_LEVEL_STEP
+    default = d.SIGNIFICANCE_LEVEL_VALUE
+    step = d.SIGNIFICANCE_LEVEL_STEP
     return st.slider(
         "Significance level",
         help=f"The level of significance of the tests. Defaults to {default:.2F}.",
@@ -45,8 +45,8 @@ def _get_significance_level():
 
 
 def _get_tau_max():
-    default = defaults.TAU_MAX_VALUE
-    step = defaults.TAU_MAX_STEP
+    default = d.TAU_MAX_VALUE
+    step = d.TAU_MAX_STEP
     return int(
         st.slider(
             "Length of test",
@@ -61,8 +61,8 @@ def _get_tau_max():
 
 
 def _get_number_of_tests():
-    default = defaults.R_VALUE
-    step = defaults.R_STEP
+    default = d.R_VALUE
+    step = d.R_STEP
     return int(
         st.slider(
             "Number of tests per concept",
@@ -80,8 +80,8 @@ def _get_number_of_tests():
 
 
 def _get_cardinality(concepts, concepts_ready):
-    default = defaults.CARDINALITY_VALUE
-    step = defaults.CARDINALITY_STEP
+    default = d.CARDINALITY_VALUE
+    step = d.CARDINALITY_STEP
     return st.slider(
         "Size of conditioning set",
         help=(
@@ -98,7 +98,7 @@ def _get_cardinality(concepts, concepts_ready):
 
 def _get_dataset_name():
     options = SUPPORTED_DATASETS
-    default_idx = options.index(defaults.DATASET_NAME)
+    default_idx = options.index(d.DATASET_NAME)
     return st.selectbox(
         "Dataset",
         options=options,
@@ -112,8 +112,8 @@ def _get_dataset_name():
 
 
 def get_model_name():
-    options = list(SUPPORTED_MODELS.keys())
-    default_idx = options.index(defaults.MODEL_NAME)
+    options = list(SUPPORTED_MODELS)
+    default_idx = options.index(d.MODEL_NAME)
     return st.selectbox(
         "Model to test",
         options=options,

app_lib/utils.py

@@ -11,16 +11,8 @@ supported_datasets_path = hf_hub_download(
     repo_type="dataset",
 )
 
-SUPPORTED_MODELS = {}
 with open(supported_models_path, "r") as f:
-    for line in f:
-        line = line.strip()
-        model_name, model_url = line.split(",")
-        SUPPORTED_MODELS[model_name] = model_url
+    SUPPORTED_MODELS = f.read().splitlines()
 
-
-SUPPORTED_DATASETS = []
 with open(supported_datasets_path, "r") as f:
-    for line in f:
-        dataset_name = line.strip()
-        SUPPORTED_DATASETS.append(dataset_name)
+    SUPPORTED_DATASETS = f.read().splitlines()

assets/results/bowl_ace.npy

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7f7746293f59199f6872a9570757b2d9d827c2733935cedba2c161181a1cc19c
+oid sha256:adaeeda897451c5548b3119fb917214b82590becfe3138158cfcf1055bcb714d
 size 226871

assets/results/gardener_ace.npy

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9e3e4a87d960a3591c6c97e10cc40a9e0727048118cc1d5670770bdd74e457ce
+oid sha256:6a7e6d291b8e7226da6af5990094501a741a99c363f04647688da6f8e71746c6
 size 226873

assets/results/gentleman_ace.npy

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a319ecde9c2299323692edf385274b6939c15a9d6c296aa70629898d8798934f
+oid sha256:fb55900dbb2870ffccc56aac8de6fde7960c3d817035bd9d461419ef0bee6b3b
 size 226874

assets/results/mathematician_ace.npy

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bf0f277e9f06b957f03c39620a11514a373754b72cf1b5e97088964d07ff7b4a
+oid sha256:ab4807601def20a04c66a2ec993a90e929172aee047e66fb8d052d8eaee438b0
 size 226873

precompute_results.py

@@ -4,7 +4,7 @@ import os
 import numpy as np
 from PIL import Image
 
-import app_lib.defaults as defaults
+from app_lib.defaults import Defaults as d
 from app_lib.test import get_testing_config, test
 
 assets_dir = "assets"
@@ -13,9 +13,9 @@ results_dir = os.path.join(assets_dir, "results")
 os.makedirs(results_dir, exist_ok=True)
 
 testing_config = get_testing_config(
-    significance_level=defaults.SIGNIFICANCE_LEVEL_VALUE,
-    tau_max=defaults.TAU_MAX_VALUE,
-    r=defaults.R_VALUE,
+    significance_level=d.SIGNIFICANCE_LEVEL_VALUE,
+    tau_max=d.TAU_MAX_VALUE,
+    r=d.R_VALUE,
 )
 
 image_presets = json.load(open(os.path.join(assets_dir, "image_presets.json")))
@@ -26,7 +26,7 @@ for _image_name, _image_presets in image_presets.items():
     _image = Image.open(_image_path)
     _class_name = _image_presets["class_name"]
     _concepts = _image_presets["concepts"]
-    _cardinality = defaults.CARDINALITY_VALUE
+    _cardinality = d.CARDINALITY_VALUE
 
     _results = test(
         testing_config,
@@ -34,8 +34,8 @@ for _image_name, _image_presets in image_presets.items():
         _class_name,
         _concepts,
         _cardinality,
-        defaults.DATASET_NAME,
-        defaults.MODEL_NAME,
+        d.DATASET_NAME,
+        d.MODEL_NAME,
         with_streamlit=False,
     )