Update

app.py

@@ -27,11 +27,6 @@ st.markdown(
             [data-testid="stHorizontalBlock"] {
                 align-items: center;
             }
-            div.stSpinner > div {
-                text-align:center;
-                align-items: center;
-                justify-content: center;
-            }
         </style>
         """,
     unsafe_allow_html=True,
@@ -44,7 +39,7 @@ st.markdown(
         Official HF Space for the paper [*I Bet You Did Not Mean That: Testing Semantci Importance via Betting*](https://arxiv.org/pdf/2405.19146), by [Jacopo Teneggi](https://jacopoteneggi.github.io) and [Jeremias Sulam](https://sites.google.com/view/jsulam).
 
         ---
-        """,
+    """,
 )
 
 if __name__ == "__main__":

app_lib/ckde.py

@@ -1,18 +1,18 @@
 import numpy as np
-import torch
 from scipy.spatial.distance import cdist
 from scipy.stats import gaussian_kde
 
+
 class cKDE:
-    def __init__(self, config, concept_class_name=None, concept_image_idx=None):
-        ckde_config = config.ckde
-        self.image_size = image_size = ckde_config.get("image_size", 128)
-        self.metric = ckde_config.get("metric", "euclidean")
-        self.scale_method = ckde_config.get("scale_method", "neff")
-        self.scale = ckde_config.get("scale", 2000)
+    def __init__(
+        self, embedding, semantics, metric="euclidean", scale_method="neff", scale=2000
+    ):
+        self.metric = metric
+        self.scale_method = scale_method
+        self.scale = scale
 
-        self.Z = self.dataset.Z
-        self.H = self.dataset.H
+        self.H = embedding
+        self.Z = semantics
 
     def _quantile_scale(self, Z_cond_dist):
         return np.quantile(Z_cond_dist, self.scale)
@@ -62,7 +62,7 @@ class cKDE:
         dist = cdist(z, self.Z, metric=self.metric)
         return np.argmin(dist, axis=-1)
 
-    def sample(self, z, cond_idx, m=1, return_images=False):
+    def sample(self, z, cond_idx, m=1):
         if z.ndim == 1:
             z = z.reshape(1, -1)
 
@@ -71,7 +71,4 @@ class cKDE:
         nn_idx = self.nearest_neighbor(sample_z)
         sample_h = self.H[nn_idx]
 
-        if return_images:
-            sample_images = torch.stack([self.dataset[_idx][0] for _idx in nn_idx])
-            return sample_z, sample_h, sample_images
         return sample_z, sample_h

app_lib/main.py

@@ -9,13 +9,7 @@ from app_lib.user_input import (
     get_image,
     get_model_name,
 )
-from app_lib.test import (
-    load_dataset,
-    load_model,
-    encode_image,
-    encode_concepts,
-    encode_class_name,
-)
+from app_lib.test import test
 
 
 def _disable():
@@ -67,27 +61,15 @@ def main(device=torch.device("cuda" if torch.cuda.is_available() else "cpu")):
             )
 
     with columns[1]:
-        if test_button:
-            with st.spinner("Loading dataset"):
-                embedding = load_dataset("imagenette", model_name)
-                time.sleep(1)
-
-            with st.spinner("Loading model"):
-                model, preprocess, tokenizer = load_model(model_name, device)
-                time.sleep(1)
-
-            with st.spinner("Encoding concepts"):
-                cbm = encode_concepts(tokenizer, model, concepts, device)
-                time.sleep(1)
-
-            with st.spinner("Preparing zero-shot classifier"):
-                classifier = encode_class_name(tokenizer, model, class_name, device)
-
-            with st.spinner("Encoding image"):
-                h = encode_image(model, preprocess, image, device)
-                z = h @ cbm.T
-                print(h.shape, cbm.shape, z.shape)
-                time.sleep(2)
-
-            st.session_state.disabled = False
-            st.experimental_rerun()
+        _, centercol, _ = st.columns(3)
+        with centercol:
+            if test_button:
+                test(
+                    image,
+                    class_name,
+                    concepts,
+                    cardinality,
+                    "imagenette",
+                    model_name,
+                    device,
+                )

app_lib/test.py

@@ -2,10 +2,29 @@ import torch
 import clip
 import open_clip
 import h5py
+import streamlit as st
+import numpy as np
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
+import ml_collections
 from huggingface_hub import hf_hub_download
 
+from ibydmt.test import xSKIT
 from app_lib.utils import SUPPORTED_MODELS
+from app_lib.ckde import cKDE
+
+rng = np.random.default_rng()
+
+testing_config = ml_collections.ConfigDict()
+testing_config.significance_level = 0.05
+testing_config.wealth = "ons"
+testing_config.bet = "tanh"
+testing_config.kernel = "rbf"
+testing_config.kernel_scale_method = "quantile"
+testing_config.kernel_scale = 0.5
+testing_config.tau_max = 200
+testing_config.r = 10
 
 
 def _get_open_clip_model(model_name, device):
@@ -39,7 +58,6 @@ def load_dataset(dataset_name, model_name):
 
 
 def load_model(model_name, device):
-    print(model_name)
     if "open_clip" in model_name:
         model, preprocess, tokenizer = _get_open_clip_model(model_name, device)
     elif "clip" in model_name:
@@ -79,6 +97,85 @@ def encode_class_name(tokenizer, model, class_name, device):
     return class_features.cpu().numpy()
 
 
+def _sample_random_subset(concept_idx, concepts, cardinality):
+    sample_idx = list(set(range(len(concepts))) - {concept_idx})
+    return rng.permutation(sample_idx)[:cardinality].tolist()
+
+
+def _test(z, concept_idx, concepts, cardinality, sampler, classifier):
+    def cond_p(z, cond_idx, m):
+        _, sample_h = sampler.sample(z, cond_idx, m=m)
+        return sample_h
+
+    def f(h):
+        output = h @ classifier.T
+        return output.squeeze()
+
+    rejected_hist, tau_hist, wealth_hist, subset_hist = [], [], [], []
+    for _ in range(testing_config.r):
+        subset_idx = _sample_random_subset(concept_idx, concepts, cardinality)
+        subset = [concepts[idx] for idx in subset_idx]
+
+        tester = xSKIT(testing_config)
+        rejected, tau = tester.test(
+            z, concept_idx, subset_idx, cond_p, f, interrupt_on_rejection=False
+        )
+        wealth = tester.wealth._wealth
+
+        rejected_hist.append(rejected)
+        tau_hist.append(tau)
+        wealth_hist.append(wealth)
+        subset_hist.append(subset)
+
+    return {
+        "concept": concepts[concept_idx],
+        "rejected": rejected_hist,
+        "tau": tau_hist,
+        "wealth": wealth_hist,
+        "subset": subset_hist,
+    }
+
+
 def test(image, class_name, concepts, cardinality, dataset_name, model_name, device):
-    model, preprocess = load_model(model_name, device)
-    print(f"loaded {model_name}")
+    with st.spinner("Loading model"):
+        model, preprocess, tokenizer = load_model(model_name, device)
+
+    with st.spinner("Encoding concepts"):
+        cbm = encode_concepts(tokenizer, model, concepts, device)
+
+    with st.spinner("Encoding image"):
+        h = encode_image(model, preprocess, image, device)
+        z = h @ cbm.T
+        z = z.squeeze()
+
+    with st.spinner("Testing"):
+        progress_bar = st.progress(0)
+
+        embedding = load_dataset("imagenette", model_name)
+        semantics = embedding @ cbm.T
+        sampler = cKDE(embedding, semantics)
+
+        classifier = encode_class_name(tokenizer, model, class_name, device)
+
+        with ThreadPoolExecutor() as executor:
+            futures = [
+                executor.submit(
+                    _test, z, concept_idx, concepts, cardinality, sampler, classifier
+                )
+                for concept_idx in range(len(concepts))
+            ]
+
+            results = []
+            for idx, future in enumerate(as_completed(futures)):
+                results.append(future.result())
+                progress_bar.progress((idx + 1) / len(concepts))
+
+        # print(results)
+        # wealth = np.empty((testing_config.tau_max, len(concepts)))
+        # wealth[:] = np.nan
+        # for _results in results:
+        #     concept_idx = concepts.index(_results["concept"])
+        #     _wealth =
+
+    st.session_state.disabled = False
+    st.experimental_rerun()

app_lib/user_input.py

@@ -80,7 +80,7 @@ def get_cardinality(concepts, concepts_ready):
         help="The number of concepts to condition model predictions on.",
         min_value=1,
         max_value=max(2, len(concepts) - 1),
-        value=1,
+        value=2,
         step=1,
         disabled=st.session_state.disabled or not concepts_ready,
     )

ibydmt/payoff.py

@@ -81,6 +81,7 @@ class KernelPayoff(Payoff):
             zip(d, null_d),
             0,
         )
+        g = g.squeeze().item()
 
         return self.bet.compute(g)
 

ibydmt/test.py

@@ -141,9 +141,12 @@ class xSKIT(SequentialTester):
         C: list[int],
         cond_p: Callable[[Float[Array, "D"], list[int], int], Float[Array, "N D2"]],
         model: Callable[[Float[Array, "N D2"]], Float[Array, "N"]],
+        interrupt_on_rejection: bool = True,
     ) -> Tuple[bool, int]:
         sample = functools.partial(self._sample, z, j, C, cond_p, model)
 
+        tau = self.tau_max - 1
+
         prev_d = np.stack(sample(), axis=1)
         for t in range(1, self.tau_max):
             y, null_y = sample()
@@ -155,5 +158,7 @@ class xSKIT(SequentialTester):
             prev_d = np.vstack([prev_d, d])
 
             if self.wealth.rejected:
-                return (True, t)
-        return (False, t)
+                tau = min(tau, t)
+                if interrupt_on_rejection:
+                    break
+        return (self.wealth.rejected, tau)

ibydmt/wealth.py

@@ -46,27 +46,30 @@ class ONS(Wealth):
     def __init__(self, config):
         super().__init__(config)
 
-        self.w = 1.0
-        self.v = 0
-        self.a = 1
+        self._w = 1.0
+        self._v = 0
+        self._a = 1
 
-        self.min_v, self.max_v = config.get("min_v", 0), config.get("max_v", 1 / 2)
-        self.wealth_flag = False
+        self._min_v, self._max_v = config.get("min_v", 0), config.get("max_v", 1 / 2)
+        self._wealth_flag = False
+
+        self._wealth = [self._w]
 
     def _update_v(self, payoff):
-        z = payoff / (1 + self.v * payoff)
-        self.a += z**2
-        self.v = max(
-            self.min_v, min(self.max_v, self.v + 2 / (2 - np.log(3)) * z / self.a)
+        z = payoff / (1 + self._v * payoff)
+        self._a += z**2
+        self._v = max(
+            self._min_v, min(self._max_v, self._v + 2 / (2 - np.log(3)) * z / self._a)
         )
 
     def update(self, payoff):
-        w = self.w * (1 + self.v * payoff)
+        w = self._w * (1 + self._v * payoff)
 
-        if w >= 0 and not self.wealth_flag:
-            self.w = w
-            if self.w >= 1 / self.significance_level:
+        if w >= 0 and not self._wealth_flag:
+            self._w = w
+            self._wealth.append(self._w)
+            if self._w >= 1 / self.significance_level:
                 self.rejected = True
             self._update_v(payoff)
         else:
-            self.wealth_flag = True
+            self._wealth_flag = True