Update

app.py

@@ -2,11 +2,6 @@ import streamlit as st
 
 from app_lib.main import main
 
-with open("style.css", "r") as f:
-    style = f.read()
-with open("header.md", "r") as f:
-    header = f.read()
-
 if "sidebar_state" not in st.session_state:
     st.session_state.sidebar_state = "collapsed"
 if "disabled" not in st.session_state:
@@ -16,9 +11,12 @@ if "results" not in st.session_state:
 
 st.set_page_config(layout="wide", initial_sidebar_state=st.session_state.sidebar_state)
 
-st.session_state.sidebar_state = "collapsed"
-st.markdown(f"<style>{style}</style>", unsafe_allow_html=True)
+with open("style.css", "r") as f:
+    style = f.read()
+with open("header.md", "r") as f:
+    header = f.read()
 
+st.markdown(f"<style>{style}</style>", unsafe_allow_html=True)
 st.markdown(header)
 
 if __name__ == "__main__":

app_lib/main.py

@@ -1,15 +1,14 @@
 import torch
 import streamlit as st
-import time
 
 from app_lib.user_input import (
-    get_cardinality,
     get_class_name,
     get_concepts,
     get_image,
     get_model_name,
+    get_advanced_settings,
 )
-from app_lib.test import test
+from app_lib.test import get_testing_config, test
 from app_lib.viz import viz_results
 
 
@@ -20,6 +19,10 @@ def _disable():
 def main(device=torch.device("cuda" if torch.cuda.is_available() else "cpu")):
     columns = st.columns([0.40, 0.60])
 
+    with columns[1]:
+        st.header("Results")
+        viz_results()
+
     with columns[0]:
         st.header("Choose Image and Concepts")
 
@@ -41,8 +44,6 @@ def main(device=torch.device("cuda" if torch.cuda.is_available() else "cpu")):
             model_name = get_model_name()
             class_name, class_ready, class_error = get_class_name()
             concepts, concepts_ready, concepts_error = get_concepts()
-            cardinality = int(len(concepts) / 2)
-            # get_cardinality(concepts, concepts_ready)
 
         ready = class_ready and concepts_ready
 
@@ -55,6 +56,10 @@ def main(device=torch.device("cuda" if torch.cuda.is_available() else "cpu")):
             st.error(error_message)
 
         with st.container():
+            significance_level, tau_max, r, cardinality = get_advanced_settings(
+                concepts, concepts_ready
+            )
+
             test_button = st.button(
                 "Test Concepts",
                 use_container_width=True,
@@ -62,25 +67,19 @@ def main(device=torch.device("cuda" if torch.cuda.is_available() else "cpu")):
                 disabled=st.session_state.disabled or not ready,
             )
 
-            with st.popover("Advanced settings", disabled=st.session_state.disabled):
-                st.markdown("Hello World 👋")
-
-    with columns[1]:
-        st.header("Results")
-
         if test_button:
             st.session_state.results = None
 
-            _, centercol, _ = st.columns(3)
-            with centercol:
-                test(
-                    image,
-                    class_name,
-                    concepts,
-                    cardinality,
-                    "imagenette",
-                    model_name,
-                    device,
-                )
-
-        viz_results()
+            testing_config = get_testing_config(
+                significance_level=significance_level, tau_max=tau_max, r=r
+            )
+            test(
+                testing_config,
+                image,
+                class_name,
+                concepts,
+                cardinality,
+                "imagenette",
+                model_name,
+                device,
+            )

app_lib/test.py

@@ -4,7 +4,6 @@ import open_clip
 import h5py
 import streamlit as st
 import numpy as np
-import pandas as pd
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
 import ml_collections
@@ -16,16 +15,6 @@ 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):
     backbone = model_name.split(":")[-1]
@@ -45,19 +34,7 @@ def _get_clip_model(model_name, device):
     return model, preprocess, tokenizer
 
 
-def load_dataset(dataset_name, model_name):
-    dataset_path = hf_hub_download(
-        repo_id="jacopoteneggi/IBYDMT",
-        filename=f"{dataset_name}_{model_name}_train.h5",
-        repo_type="dataset",
-    )
-
-    with h5py.File(dataset_path, "r") as dataset:
-        embedding = dataset["embedding"][:]
-    return embedding
-
-
-def load_model(model_name, device):
+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:
@@ -67,7 +44,7 @@ def load_model(model_name, device):
 
 @torch.no_grad()
 @torch.cuda.amp.autocast()
-def encode_concepts(tokenizer, model, concepts, device):
+def _encode_concepts(tokenizer, model, concepts, device):
     concepts_text = tokenizer(concepts).to(device)
 
     concept_features = model.encode_text(concepts_text)
@@ -77,7 +54,7 @@ def encode_concepts(tokenizer, model, concepts, device):
 
 @torch.no_grad()
 @torch.cuda.amp.autocast()
-def encode_image(model, preprocess, image, device):
+def _encode_image(model, preprocess, image, device):
     image = preprocess(image)
     image = image.unsqueeze(0)
     image = image.to(device)
@@ -89,7 +66,7 @@ def encode_image(model, preprocess, image, device):
 
 @torch.no_grad()
 @torch.cuda.amp.autocast()
-def encode_class_name(tokenizer, model, class_name, device):
+def _encode_class_name(tokenizer, model, class_name, device):
     class_text = tokenizer([f"A photo of a {class_name}"]).to(device)
 
     class_features = model.encode_text(class_text)
@@ -97,12 +74,24 @@ def encode_class_name(tokenizer, model, class_name, device):
     return class_features.cpu().numpy()
 
 
+def _load_dataset(dataset_name, model_name):
+    dataset_path = hf_hub_download(
+        repo_id="jacopoteneggi/IBYDMT",
+        filename=f"{dataset_name}_{model_name}_train.h5",
+        repo_type="dataset",
+    )
+
+    with h5py.File(dataset_path, "r") as dataset:
+        embedding = dataset["embedding"][:]
+    return embedding
+
+
 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 _test(testing_config, 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
@@ -118,9 +107,16 @@ def _test(z, concept_idx, concepts, cardinality, sampler, classifier):
 
         tester = xSKIT(testing_config)
         rejected, tau = tester.test(
-            z, concept_idx, subset_idx, cond_p, f, interrupt_on_rejection=False
+            z,
+            concept_idx,
+            subset_idx,
+            cond_p,
+            f,
+            interrupt_on="max_wealth",
+            max_wealth=100,
         )
         wealth = tester.wealth._wealth
+        wealth = wealth + [wealth[-1]] * (testing_config.tau_max - len(wealth))
 
         rejected_hist.append(rejected)
         tau_hist.append(tau)
@@ -136,60 +132,99 @@ def _test(z, concept_idx, concepts, cardinality, sampler, classifier):
     }
 
 
-def test(image, class_name, concepts, cardinality, dataset_name, model_name, device):
+def get_testing_config(**kwargs):
+    testing_config = st.session_state.testing_config = ml_collections.ConfigDict()
+    testing_config.significance_level = kwargs.get("significance_level", 0.05)
+    testing_config.wealth = kwargs.get("wealth", "ons")
+    testing_config.bet = kwargs.get("bet", "tanh")
+    testing_config.kernel = kwargs.get("kernel", "rbf")
+    testing_config.kernel_scale_method = kwargs.get("kernel_scale_method", "quantile")
+    testing_config.kernel_scale = kwargs.get("kernel_scale", 0.5)
+    testing_config.tau_max = kwargs.get("tau_max", 200)
+    testing_config.r = kwargs.get("r", 10)
+    return testing_config
+
+
+def test(
+    testing_config,
+    image,
+    class_name,
+    concepts,
+    cardinality,
+    dataset_name,
+    model_name,
+    device,
+):
     with st.spinner("Loading model"):
-        model, preprocess, tokenizer = load_model(model_name, device)
+        model, preprocess, tokenizer = _load_model(model_name, device)
 
     with st.spinner("Encoding concepts"):
-        cbm = encode_concepts(tokenizer, model, concepts, device)
+        cbm = _encode_concepts(tokenizer, model, concepts, device)
 
     with st.spinner("Encoding image"):
-        h = encode_image(model, preprocess, image, device)
+        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))
-
-        rejected = np.empty((testing_config.r, len(concepts)))
-        tau = np.empty((testing_config.r, len(concepts)))
-        wealth = np.empty((testing_config.r, testing_config.tau_max, len(concepts)))
+    progress_bar = st.progress(
+        0,
+        text=f"Testing concepts (can take a few minutes) [0 / {len(concepts)} completed]",
+    )
 
-        for _results in results:
-            concept_idx = concepts.index(_results["concept"])
+    embedding = _load_dataset(dataset_name, model_name)
+    semantics = embedding @ cbm.T
+    sampler = cKDE(embedding, semantics)
 
-            rejected[:, concept_idx] = np.array(_results["rejected"])
-            tau[:, concept_idx] = np.array(_results["tau"])
-            wealth[:, :, concept_idx] = np.array(_results["wealth"])
+    classifier = _encode_class_name(tokenizer, model, class_name, device)
 
-        tau /= testing_config.tau_max
+    progress_bar.progress(
+        1 / (len(concepts) + 1),
+        text=f"Testing concepts (can take a few minutes) [0 / {len(concepts)} completed]",
+    )
 
-        st.session_state.results = {
-            "significance_level": testing_config.significance_level,
-            "concepts": concepts,
-            "rejected": rejected,
-            "tau": tau,
-            "wealth": wealth,
-        }
+    with ThreadPoolExecutor() as executor:
+        futures = [
+            executor.submit(
+                _test,
+                testing_config,
+                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 + 2) / (len(concepts) + 1),
+                text=f"Testing concepts (can take a few minutes) [{idx + 1} / {len(concepts)} completed]",
+            )
+
+    rejected = np.empty((testing_config.r, len(concepts)))
+    tau = np.empty((testing_config.r, len(concepts)))
+    wealth = np.empty((testing_config.r, testing_config.tau_max, len(concepts)))
+
+    for _results in results:
+        concept_idx = concepts.index(_results["concept"])
+
+        rejected[:, concept_idx] = np.array(_results["rejected"])
+        tau[:, concept_idx] = np.array(_results["tau"])
+        wealth[:, :, concept_idx] = np.array(_results["wealth"])
+
+    tau /= testing_config.tau_max
+
+    st.session_state.results = {
+        "significance_level": testing_config.significance_level,
+        "concepts": concepts,
+        "rejected": rejected,
+        "tau": tau,
+        "wealth": wealth,
+    }
 
     st.session_state.disabled = False
     st.experimental_rerun()

app_lib/user_input.py

@@ -20,6 +20,82 @@ def _validate_concepts(concepts):
         return (False, "Maximum 10 concepts allowed")
     return (True, None)
 
+
+def _get_significance_level():
+    DEFAULT = 0.05
+    return st.slider(
+        "Significance level",
+        help=" ".join(
+            [
+                "The level of significance of the tests.",
+                f"Defaults to {DEFAULT:.2F}.",
+            ]
+        ),
+        min_value=0.01,
+        max_value=1.0,
+        value=DEFAULT,
+        step=0.01,
+        disabled=st.session_state.disabled,
+    )
+
+
+def _get_tau_max():
+    DEFAULT = 200
+    return int(
+        st.slider(
+            "Duration of test",
+            help=" ".join(
+                [
+                    "The maximum number of steps for each test.",
+                    f"Defaults to {DEFAULT}.",
+                ]
+            ),
+            min_value=1,
+            max_value=1000,
+            step=1,
+            value=DEFAULT,
+            disabled=st.session_state.disabled,
+        )
+    )
+
+
+def _get_number_of_tests():
+    DEFAULT = 20
+    return int(
+        st.slider(
+            "Number of tests per concept",
+            help=" ".join(
+                [
+                    "The number of tests to average for each concept.",
+                    f"Defaults to {DEFAULT}.",
+                ]
+            ),
+            min_value=1,
+            max_value=100,
+            step=1,
+            value=DEFAULT,
+            disabled=st.session_state.disabled,
+        )
+    )
+
+
+def _get_cardinality(concepts, concepts_ready):
+    return st.slider(
+        "Size of conditioning set",
+        help=" ".join(
+            [
+                "The number of concepts to condition model predictions on.",
+                "Defaults to half of the number of concepts.",
+            ]
+        ),
+        min_value=1,
+        max_value=max(2, len(concepts) - 1),
+        value=int(len(concepts) / 2),
+        step=1,
+        disabled=st.session_state.disabled or not concepts_ready,
+    )
+
+
 def get_model_name():
     return st.selectbox(
         "Model to test",
@@ -58,8 +134,8 @@ def get_class_name():
 
 def get_concepts():
     concepts = st.text_area(
-        "Concepts to test (max 10)",
-        help="List of concepts to test the predictions of the model with. Write one concept per line.",
+        "Concepts to test",
+        help="List of concepts to test the predictions of the model with. Write one concept per line. Maximum 10 concepts allowed.",
         height=160,
         value="piano\ncute\nwhiskers\nmusic\nwild",
         disabled=st.session_state.disabled,
@@ -73,13 +149,11 @@ def get_concepts():
     return concepts, concepts_ready, concepts_error
 
 
-def get_cardinality(concepts, concepts_ready):
-    return st.slider(
-        "Size of conditioning set",
-        help="The number of concepts to condition model predictions on.",
-        min_value=1,
-        max_value=max(2, len(concepts) - 1),
-        value=2,
-        step=1,
-        disabled=st.session_state.disabled or not concepts_ready,
-    )
+def get_advanced_settings(concepts, concepts_ready):
+    with st.popover("Advanced settings", disabled=st.session_state.disabled):
+        significance_level = _get_significance_level()
+        tau_max = _get_tau_max()
+        r = _get_number_of_tests()
+        cardinality = _get_cardinality(concepts, concepts_ready)
+
+    return significance_level, tau_max, r, cardinality

header.md

@@ -1,5 +1,3 @@
 # 🤔 I Bet You Did Not Mean That
 
-Official HF Space for the paper [*I Bet You Did Not Mean That: Testing Semantic 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).
-
----
+Official 🤗 Space for the paper [*I Bet You Did Not Mean That: Testing Semantic 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).

ibydmt/test.py

@@ -141,7 +141,8 @@ 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,
+        interrupt_on: str = "rejection",
+        max_wealth: float = None,
     ) -> Tuple[bool, int]:
         sample = functools.partial(self._sample, z, j, C, cond_p, model)
 
@@ -159,6 +160,8 @@ class xSKIT(SequentialTester):
 
             if self.wealth.rejected:
                 tau = min(tau, t)
-                if interrupt_on_rejection:
+                if interrupt_on == "rejection":
+                    break
+                if interrupt_on == "max_wealth" and self.wealth._w >= max_wealth:
                     break
         return (self.wealth.rejected, tau)

requirements.txt

@@ -4,4 +4,5 @@ open_clip_torch
 h5py
 ml_collections
 jaxtyping
-scikit-learn
+scikit-learn
+plotly

style.css

@@ -33,7 +33,22 @@ h1 {
         }
     }
 
+    button:active {
+        background: white;
+    }
+
+    button:focus:not(:active) {
+        color: rgb(49, 51, 63);
+    }
+
     button:hover>div:first-of-type>p {
         text-decoration: underline;
     }
+}
+
+[data-testid="stSpinner"] {
+    >div {
+        display: flex;
+        justify-content: center;
+    }
 }