Upload 126 files

.gitattributes

@@ -34,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 DFA-MoE/docs/intro.png filter=lfs diff=lfs merge=lfs -text
+docs/intro.png filter=lfs diff=lfs merge=lfs -text

MTIL_datasets/__init__.py

@@ -0,0 +1,4 @@
+# This dataset is originally proposed by Zangwei Zheng et al. (https://arxiv.org/abs/2303.06628)
+# Code here is based on CoOp and its following works. (https://github.com/KaiyangZhou/CoOp)
+# Modified by Longxiang Tang (lloong.x@gmail.com) to release the dependence of Dassl lib.
+# To prepare data, please refer to https://github.com/muzairkhattak/PromptSRC/blob/main/docs/DATASETS.md

MTIL_datasets/caltech101.py

@@ -0,0 +1,92 @@
+import os
+import pickle
+
+from .utils import *
+
+from .oxford_pets import OxfordPets
+from .dtd import DescribableTextures as DTD
+
+IGNORED = ["BACKGROUND_Google", "Faces_easy"]
+NEW_CNAMES = {
+    "airplanes": "airplane",
+    "Faces": "face",
+    "Leopards": "leopard",
+    "Motorbikes": "motorbike",
+}
+
+
+class Caltech101(DatasetBase):
+
+    dataset_dir = "caltech-101"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.image_dir = os.path.join(self.dataset_dir, "101_ObjectCategories")
+        self.split_path = os.path.join(self.dataset_dir, "split_zhou_Caltech101.json")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        if os.path.exists(self.split_path):
+            train, val, test = OxfordPets.read_split(self.split_path, self.image_dir)
+        else:
+            train, val, test = DTD.read_and_split_data(self.image_dir, ignored=IGNORED, new_cnames=NEW_CNAMES)
+            OxfordPets.save_split(train, val, test, self.split_path, self.image_dir)
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        subsample = subsample_classes
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample)
+
+        self.templates = [
+            lambda c: f"a photo of a {c}.",
+            lambda c: f"a painting of a {c}.",
+            lambda c: f"a plastic {c}.",
+            lambda c: f"a sculpture of a {c}.",
+            lambda c: f"a sketch of a {c}.",
+            lambda c: f"a tattoo of a {c}.",
+            lambda c: f"a toy {c}.",
+            lambda c: f"a rendition of a {c}.",
+            lambda c: f"a embroidered {c}.",
+            lambda c: f"a cartoon {c}.",
+            lambda c: f"a {c} in a video game.",
+            lambda c: f"a plushie {c}.",
+            lambda c: f"a origami {c}.",
+            lambda c: f"art of a {c}.",
+            lambda c: f"graffiti of a {c}.",
+            lambda c: f"a drawing of a {c}.",
+            lambda c: f"a doodle of a {c}.",
+            lambda c: f"a photo of the {c}.",
+            lambda c: f"a painting of the {c}.",
+            lambda c: f"the plastic {c}.",
+            lambda c: f"a sculpture of the {c}.",
+            lambda c: f"a sketch of the {c}.",
+            lambda c: f"a tattoo of the {c}.",
+            lambda c: f"the toy {c}.",
+            lambda c: f"a rendition of the {c}.",
+            lambda c: f"the embroidered {c}.",
+            lambda c: f"the cartoon {c}.",
+            lambda c: f"the {c} in a video game.",
+            lambda c: f"the plushie {c}.",
+            lambda c: f"the origami {c}.",
+            lambda c: f"art of the {c}.",
+            lambda c: f"graffiti of the {c}.",
+            lambda c: f"a drawing of the {c}.",
+            lambda c: f"a doodle of the {c}.",
+        ]
+
+        super().__init__(train_x=train, val=val, test=test)

MTIL_datasets/cifar10.py

@@ -0,0 +1,94 @@
+import os
+from typing import List
+
+from .utils import *  # Datum, DatasetBase
+from .oxford_pets import OxfordPets
+
+try:
+    from torchvision.datasets import CIFAR10 as TorchCIFAR10
+except Exception as e:
+    TorchCIFAR10 = None
+    print(f"Warning: torchvision not available for CIFAR10: {e}")
+
+
+CIFAR10_CLASSES: List[str] = [
+    'airplane',
+    'automobile',
+    'bird',
+    'cat',
+    'deer',
+    'dog',
+    'frog',
+    'horse',
+    'ship',
+    'truck',
+]
+
+CIFAR10_TEMPLATES: List[str] = [
+    'a photo of a {}.',
+    'a blurry photo of a {}.',
+    'a black and white photo of a {}.',
+    'a low contrast photo of a {}.',
+    'a high contrast photo of a {}.',
+    'a bad photo of a {}.',
+    'a good photo of a {}.',
+    'a photo of a small {}.',
+    'a photo of a big {}.',
+    'a photo of the {}.',
+    'a blurry photo of the {}.',
+    'a black and white photo of the {}.',
+    'a low contrast photo of the {}.',
+    'a high contrast photo of the {}.',
+    'a bad photo of the {}.',
+    'a good photo of the {}.',
+    'a photo of the small {}.',
+    'a photo of the big {}.',
+]
+
+
+class CIFAR10(DatasetBase):
+
+    dataset_dir = "cifar10"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        if TorchCIFAR10 is None:
+            raise ImportError("torchvision is required for CIFAR10 dataset. Please install torchvision.")
+
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+
+        # Use torchvision to download/load data to dataset_dir
+        train_ds = TorchCIFAR10(root=self.dataset_dir, train=True, download=True)
+        test_ds = TorchCIFAR10(root=self.dataset_dir, train=False, download=True)
+
+        # Build Datum lists
+        trainval = []
+        # train_ds.data: numpy array HWC, train_ds.targets: list[int]
+        for idx in range(len(train_ds.data)):
+            img = Image.fromarray(train_ds.data[idx])
+            label = int(train_ds.targets[idx])
+            classname = CIFAR10_CLASSES[label]
+            trainval.append(Datum(impath=img, label=label, classname=classname))
+
+        test = []
+        for idx in range(len(test_ds.data)):
+            img = Image.fromarray(test_ds.data[idx])
+            label = int(test_ds.targets[idx])
+            classname = CIFAR10_CLASSES[label]
+            test.append(Datum(impath=img, label=label, classname=classname))
+
+        # Split train/val
+        train, val = OxfordPets.split_trainval(trainval)
+
+        # Few-shot sampling if requested
+        if num_shots >= 1:
+            train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+            val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+
+        # Optional class subsampling (base/new)
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample_classes)
+
+        # Templates
+        self.templates = CIFAR10_TEMPLATES
+
+        super().__init__(train_x=train, val=val, test=test)

MTIL_datasets/cifar100.py

@@ -0,0 +1,89 @@
+import os
+import pickle
+import numpy as np
+
+from .utils import *
+
+from .oxford_pets import OxfordPets
+from .dtd import DescribableTextures as DTD
+
+
+class CIFAR100(DatasetBase):
+
+    dataset_dir = "cifar100"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+
+        file_path = os.path.join(root, self.dataset_dir, 'train')
+        with open(file_path, "rb") as f:
+            entry = pickle.load(f, encoding="latin1")
+            trainval_data = entry["data"]
+            if "labels" in entry:
+                trainval_targets = entry["labels"]
+            else:
+                trainval_targets = entry["fine_labels"]
+        trainval_data = trainval_data.reshape(-1, 3, 32, 32)
+        trainval_data = trainval_data.transpose((0, 2, 3, 1))
+
+        file_path = os.path.join(root, self.dataset_dir, 'test')
+        with open(file_path, "rb") as f:
+            entry = pickle.load(f, encoding="latin1")
+            test_data = entry["data"]
+            if "labels" in entry:
+                test_targets = entry["labels"]
+            else:
+                test_targets = entry["fine_labels"]
+        test_data = test_data.reshape(-1, 3, 32, 32)
+        test_data = test_data.transpose((0, 2, 3, 1))
+
+        path = os.path.join(self.dataset_dir, "meta")
+        with open(path, "rb") as infile:
+            data = pickle.load(infile, encoding="latin1")
+            classes = data["fine_label_names"]
+        classes = [s.replace("_", " ") for s in classes]
+
+        trainval = []
+        for idx in range(trainval_data.shape[0]):
+            item = Datum(impath=Image.fromarray(trainval_data[idx]), 
+                         label=int(trainval_targets[idx]), classname=classes[trainval_targets[idx]])
+            trainval.append(item)
+        
+        test = []
+        for idx in range(test_data.shape[0]):
+            item = Datum(impath=Image.fromarray(test_data[idx]), 
+                         label=int(test_targets[idx]), classname=classes[test_targets[idx]])
+            test.append(item)
+        
+        train, val = OxfordPets.split_trainval(trainval)
+        
+        if num_shots >= 1:
+            train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+            val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+
+        subsample = subsample_classes
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample)
+
+        self.templates = [
+            lambda c : f'a bad photo of a {c}.',
+            lambda c : f'a blurry photo of a {c}.',
+            lambda c : f'a black and white photo of a {c}.',
+            lambda c : f'a low contrast photo of a {c}.',
+            lambda c : f'a high contrast photo of a {c}.',
+            lambda c : f'a photo of a {c}.',
+            lambda c : f'a good photo of a {c}.',
+            lambda c : f'a photo of a small {c}.',
+            lambda c : f'a photo of a big {c}.',
+            lambda c : f'a photo of the {c}.',
+            lambda c : f'a blurry photo of the {c}.',
+            lambda c : f'a black and white photo of the {c}.',
+            lambda c : f'a low contrast photo of the {c}.',
+            lambda c : f'a high contrast photo of the {c}.',
+            lambda c : f'a bad photo of the {c}.',
+            lambda c : f'a good photo of the {c}.',
+            lambda c : f'a photo of the small {c}.',
+            lambda c : f'a photo of the big {c}.',
+        ]
+
+        super().__init__(train_x=train, val=val, test=test)

MTIL_datasets/clevr_count.py

@@ -0,0 +1,118 @@
+import os
+import pickle
+from typing import List
+
+from .utils import Datum, DatasetBase, mkdir_if_missing
+from .oxford_pets import OxfordPets
+from .utils import read_json, write_json
+
+# Class list and templates as specified
+CLEVR_COUNT_CLASSES: List[str] = [
+    '10', '3', '4', '5', '6', '7', '8', '9'
+]
+
+CLEVR_COUNT_TEMPLATES: List[str] = [
+    'a photo of {} objects.',
+]
+
+
+class CLEVRCount(DatasetBase):
+
+    dataset_dir = 'clevr'
+
+    def __init__(self, root, num_shots: int = 0, seed: int = 1, subsample_classes: str = 'all'):
+        # Root and directories
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.images_dir = os.path.join(self.dataset_dir, 'images')
+        self.scenes_dir = os.path.join(self.dataset_dir, 'scenes')
+        self.split_path = os.path.join(self.dataset_dir, 'split_custom_CLEVRCount.json')
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, 'split_fewshot')
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        # Required files
+        train_scenes = os.path.join(self.scenes_dir, 'CLEVR_train_scenes.json')
+        val_scenes = os.path.join(self.scenes_dir, 'CLEVR_val_scenes.json')
+        if not os.path.isfile(train_scenes) or not os.path.isfile(val_scenes):
+            raise FileNotFoundError(
+                f"CLEVRCount expects scenes JSON at {train_scenes} and {val_scenes}"
+            )
+
+        # Load or build split
+        if os.path.exists(self.split_path):
+            train, val, test = OxfordPets.read_split(self.split_path, self.dataset_dir)
+        else:
+            trainval = self._read_scenes(train_scenes, split='train')
+            test = self._read_scenes(val_scenes, split='val')
+            train, val = OxfordPets.split_trainval(trainval)
+            OxfordPets.save_split(train, val, test, self.split_path, self.dataset_dir)
+
+        # Few-shot
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, 'rb') as f:
+                    data = pickle.load(f)
+                    train, val = data['train'], data['val']
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {'train': train, 'val': val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, 'wb') as f:
+                    pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
+
+        # Optional class subsampling (base/new)
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample_classes)
+
+        # Templates
+        self.templates = CLEVR_COUNT_TEMPLATES
+
+        # Debug stats
+        try:
+            def _hist(items: List[Datum]):
+                from collections import Counter
+                cnt = Counter([it.label for it in items])
+                out = {i: int(cnt.get(i, 0)) for i in range(len(CLEVR_COUNT_CLASSES))}
+                return out
+        except Exception as e:
+            print(f"CLEVRCount stats printing failed: {e}")
+
+        super().__init__(train_x=train, val=val, test=test)
+        # Ensure class metadata is stable and matches the provided list
+        self._classnames = CLEVR_COUNT_CLASSES
+        self._lab2cname = {i: c for i, c in enumerate(CLEVR_COUNT_CLASSES)}
+        self._num_classes = len(CLEVR_COUNT_CLASSES)
+
+    def _read_scenes(self, json_path: str, split: str) -> List[Datum]:
+        """Read CLEVR scenes JSON and construct a list of Datum entries.
+        Only images whose object count appears in CLEVR_COUNT_CLASSES are kept.
+        """
+        obj = read_json(json_path)
+        scenes = obj.get('scenes', [])
+        # Map count -> class index
+        lab2idx = {int(c): i for i, c in enumerate(CLEVR_COUNT_CLASSES)}
+        items: List[Datum] = []
+        for sc in scenes:
+            # Some files use key 'split', some typos list 'spit'; be robust
+            image_filename = sc.get('image_filename', None)
+            if not image_filename:
+                continue
+            n_objects = sc.get('objects', [])
+            try:
+                num = int(len(n_objects))
+            except Exception:
+                continue
+            if num not in lab2idx:
+                # skip counts not in the configured class list
+                continue
+            label_i = lab2idx[num]
+            # Build absolute path to the image based on split
+            if split not in ['train', 'val', 'test']:
+                split_dir = 'train'
+            else:
+                split_dir = split
+            impath = os.path.join(self.dataset_dir, 'images', split_dir, image_filename)
+            items.append(Datum(impath=impath, label=label_i, classname=str(num)))
+        return items

MTIL_datasets/country211.py

@@ -0,0 +1,386 @@
+import os
+import json
+import pickle
+import re
+
+from .utils import *
+from .oxford_pets import OxfordPets
+from .dtd import DescribableTextures as DTD
+
+
+class Country211(DatasetBase):
+    """
+    Country211 dataset loader for MTIL.
+    Expected structure:
+      <root>/country211/
+        ├─ train/<ISO2>/*.jpg
+        ├─ valid/<ISO2>/*.jpg
+        └─ test/<ISO2>/*.jpg
+
+    Where <ISO2> are ISO-3166 alpha-2 country codes (e.g., AD, US, CN).
+    Class names are mapped from ISO2 codes to provided human-readable names.
+    Optionally, a JSON mapping file can override defaults:
+      <root>/country211/iso2_to_name.json  => {"AD": "Andorra", ...}
+    """
+
+    dataset_dir = "country211"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        # Prefer explicit split directories if they exist
+        train_dir = os.path.join(self.dataset_dir, "train")
+        valid_dir = os.path.join(self.dataset_dir, "valid")
+        test_dir = os.path.join(self.dataset_dir, "test")
+
+        use_folder_splits = os.path.isdir(train_dir) and os.path.isdir(valid_dir) and os.path.isdir(test_dir)
+
+        if use_folder_splits:
+            # Build ISO2 -> name mapping (allow override via json)
+            iso_map = self._load_iso2_to_name()
+            # Create unified code set across splits to keep label mapping stable
+            all_codes = set()
+            for d in [train_dir, valid_dir, test_dir]:
+                if os.path.isdir(d):
+                    for code in listdir_nohidden(d):
+                        code_path = os.path.join(d, code)
+                        if os.path.isdir(code_path):
+                            all_codes.add(code)
+            # Validate code format and mapping coverage early
+            override_path = os.path.join(self.dataset_dir, 'iso2_to_name.json')
+            pat = re.compile(r'^[A-Z]{2}$')
+            invalid_codes = sorted([c for c in all_codes if not (pat.match(c) or c == 'XK')])
+            if invalid_codes:
+                raise ValueError(
+                    "Country211: Found invalid ISO2 code folder names: {}. "
+                    "Codes must be two uppercase letters (e.g., 'US', 'CN') or 'XK'. "
+                    "Please rename these folders accordingly.".format(invalid_codes)
+                )
+            unknown_codes = sorted([c for c in all_codes if c not in iso_map])
+            if unknown_codes:
+                raise ValueError(
+                    "Country211: ISO2 codes missing from mapping: {}. "
+                    "Add them to {} as a JSON dict, e.g., {\"XX\": \"Country Name\"}.".format(
+                        unknown_codes, override_path
+                    )
+                )
+            # Sort by human name (fallback to code) for stable labels
+            def _code_to_name(c):
+                return iso_map.get(c, c)
+            sorted_codes = sorted(list(all_codes), key=lambda c: _code_to_name(c))
+            code_to_label = {c: i for i, c in enumerate(sorted_codes)}
+
+            train = self._read_split_dir(train_dir, code_to_label, iso_map)
+            val = self._read_split_dir(valid_dir, code_to_label, iso_map)
+            test = self._read_split_dir(test_dir, code_to_label, iso_map)
+        else:
+            # Fallbacks: JSON split or naive folder split
+            image_dir = os.path.join(self.dataset_dir, "images")
+            self.image_dir = image_dir if os.path.isdir(image_dir) else self.dataset_dir
+            split_path = os.path.join(self.dataset_dir, "split_zhou_Country211.json")
+            if os.path.exists(split_path):
+                train, val, test = OxfordPets.read_split(split_path, self.image_dir)
+            else:
+                train, val, test = DTD.read_and_split_data(self.image_dir)
+                OxfordPets.save_split(train, val, test, split_path, self.image_dir)
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        subsample = subsample_classes
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample)
+
+        # Templates provided by user
+        self.templates = [
+            lambda c: f'a photo i took in {c}.',
+            lambda c: f'a photo i took while visiting {c}.',
+            lambda c: f'a photo from my home country of {c}.',
+            lambda c: f'a photo from my visit to {c}.',
+            lambda c: f'a photo showing the country of {c}.',
+        ]
+
+        super().__init__(train_x=train, val=val, test=test)
+
+    def _read_split_dir(self, split_dir, code_to_label, iso_map):
+        items = []
+        if not os.path.isdir(split_dir):
+            return items
+        codes = listdir_nohidden(split_dir)
+        for code in codes:
+            class_dir = os.path.join(split_dir, code)
+            if not os.path.isdir(class_dir):
+                continue
+            label = code_to_label.get(code)
+            if label is None:
+                raise ValueError(
+                    f"Country211: Inconsistent label mapping for code '{code}' in split '{split_dir}'. "
+                    f"This indicates an internal mismatch between discovered codes and label map."
+                )
+            try:
+                cname = iso_map[code]
+            except KeyError:
+                raise ValueError(
+                    f"Country211: Missing ISO mapping for code '{code}'. "
+                    f"Please add it to '{os.path.join(self.dataset_dir, 'iso2_to_name.json')}'."
+                )
+            for fname in listdir_nohidden(class_dir):
+                impath = os.path.join(class_dir, fname)
+                items.append(Datum(impath=impath, label=label, classname=cname))
+        return items
+
+    def _load_iso2_to_name(self):
+        """Load ISO2->country name mapping, allow local JSON override.
+        Fallback to built-in mapping; unknown codes map to themselves at usage time.
+        """
+        override_path = os.path.join(self.dataset_dir, 'iso2_to_name.json')
+        if os.path.exists(override_path):
+            try:
+                data = read_json(override_path)
+                if not isinstance(data, dict):
+                    raise ValueError(
+                        f"Country211: Expected a JSON object in '{override_path}', got {type(data)}"
+                    )
+                # Validate keys as ISO2 (two uppercase letters) or 'XK', and values are non-empty strings
+                pat = re.compile(r'^[A-Z]{2}$')
+                bad_keys = [k for k in data.keys() if not (isinstance(k, str) and (pat.match(k) or k == 'XK'))]
+                bad_vals = [k for k, v in data.items() if not (isinstance(v, str) and v.strip())]
+                if bad_keys or bad_vals:
+                    raise ValueError(
+                        (
+                            "Country211: Invalid iso2_to_name.json entries. "
+                            f"Bad keys (must be ISO2 like 'US'): {bad_keys}. "
+                            f"Bad values (must be non-empty strings) for keys: {bad_vals}."
+                        )
+                    )
+                return data
+            except Exception as e:
+                raise ValueError(f"Country211: Failed to load '{override_path}': {e}")
+        # Built-in mapping aligned with provided classes
+        return {
+            'AD': 'Andorra',
+            'AE': 'United Arab Emirates',
+            'AF': 'Afghanistan',
+            'AG': 'Antigua and Barbuda',
+            'AI': 'Anguilla',
+            'AL': 'Albania',
+            'AM': 'Armenia',
+            'AO': 'Angola',
+            'AQ': 'Antarctica',
+            'AR': 'Argentina',
+            'AT': 'Austria',
+            'AU': 'Australia',
+            'AW': 'Aruba',
+            'AX': 'Aland Islands',
+            'AZ': 'Azerbaijan',
+            'BA': 'Bosnia and Herzegovina',
+            'BB': 'Barbados',
+            'BD': 'Bangladesh',
+            'BE': 'Belgium',
+            'BF': 'Burkina Faso',
+            'BG': 'Bulgaria',
+            'BH': 'Bahrain',
+            'BJ': 'Benin',
+            'BM': 'Bermuda',
+            'BN': 'Brunei Darussalam',
+            'BO': 'Bolivia',
+            'BQ': 'Bonaire, Saint Eustatius and Saba',
+            'BR': 'Brazil',
+            'BS': 'Bahamas',
+            'BT': 'Bhutan',
+            'BW': 'Botswana',
+            'BY': 'Belarus',
+            'BZ': 'Belize',
+            'CA': 'Canada',
+            'CD': 'DR Congo',
+            'CF': 'Central African Republic',
+            'CH': 'Switzerland',
+            'CI': "Cote d'Ivoire",
+            'CK': 'Cook Islands',
+            'CL': 'Chile',
+            'CM': 'Cameroon',
+            'CN': 'China',
+            'CO': 'Colombia',
+            'CR': 'Costa Rica',
+            'CU': 'Cuba',
+            'CV': 'Cabo Verde',
+            'CW': 'Curacao',
+            'CY': 'Cyprus',
+            'CZ': 'Czech Republic',
+            'DE': 'Germany',
+            'DK': 'Denmark',
+            'DM': 'Dominica',
+            'DO': 'Dominican Republic',
+            'DZ': 'Algeria',
+            'EC': 'Ecuador',
+            'EE': 'Estonia',
+            'EG': 'Egypt',
+            'ES': 'Spain',
+            'ET': 'Ethiopia',
+            'FI': 'Finland',
+            'FJ': 'Fiji',
+            'FK': 'Falkland Islands',
+            'FO': 'Faeroe Islands',
+            'FR': 'France',
+            'GA': 'Gabon',
+            'GB': 'United Kingdom',
+            'GD': 'Grenada',
+            'GE': 'Georgia',
+            'GF': 'French Guiana',
+            'GG': 'Guernsey',
+            'GH': 'Ghana',
+            'GI': 'Gibraltar',
+            'GL': 'Greenland',
+            'GM': 'Gambia',
+            'GP': 'Guadeloupe',
+            'GR': 'Greece',
+            'GS': 'South Georgia and South Sandwich Is.',
+            'GT': 'Guatemala',
+            'GU': 'Guam',
+            'GY': 'Guyana',
+            'HK': 'Hong Kong',
+            'HN': 'Honduras',
+            'HR': 'Croatia',
+            'HT': 'Haiti',
+            'HU': 'Hungary',
+            'ID': 'Indonesia',
+            'IE': 'Ireland',
+            'IL': 'Israel',
+            'IM': 'Isle of Man',
+            'IN': 'India',
+            'IQ': 'Iraq',
+            'IR': 'Iran',
+            'IS': 'Iceland',
+            'IT': 'Italy',
+            'JE': 'Jersey',
+            'JM': 'Jamaica',
+            'JO': 'Jordan',
+            'JP': 'Japan',
+            'KE': 'Kenya',
+            'KG': 'Kyrgyz Republic',
+            'KH': 'Cambodia',
+            'KN': 'St. Kitts and Nevis',
+            'KP': 'North Korea',
+            'KR': 'South Korea',
+            'KW': 'Kuwait',
+            'KY': 'Cayman Islands',
+            'KZ': 'Kazakhstan',
+            'LA': 'Laos',
+            'LB': 'Lebanon',
+            'LC': 'St. Lucia',
+            'LI': 'Liechtenstein',
+            'LK': 'Sri Lanka',
+            'LR': 'Liberia',
+            'LT': 'Lithuania',
+            'LU': 'Luxembourg',
+            'LV': 'Latvia',
+            'LY': 'Libya',
+            'MA': 'Morocco',
+            'MC': 'Monaco',
+            'MD': 'Moldova',
+            'ME': 'Montenegro',
+            'MF': 'Saint-Martin',
+            'MG': 'Madagascar',
+            'MK': 'Macedonia',
+            'ML': 'Mali',
+            'MM': 'Myanmar',
+            'MN': 'Mongolia',
+            'MO': 'Macau',
+            'MQ': 'Martinique',
+            'MR': 'Mauritania',
+            'MT': 'Malta',
+            'MU': 'Mauritius',
+            'MV': 'Maldives',
+            'MW': 'Malawi',
+            'MX': 'Mexico',
+            'MY': 'Malaysia',
+            'MZ': 'Mozambique',
+            'NA': 'Namibia',
+            'NC': 'New Caledonia',
+            'NG': 'Nigeria',
+            'NI': 'Nicaragua',
+            'NL': 'Netherlands',
+            'NO': 'Norway',
+            'NP': 'Nepal',
+            'NZ': 'New Zealand',
+            'OM': 'Oman',
+            'PA': 'Panama',
+            'PE': 'Peru',
+            'PF': 'French Polynesia',
+            'PG': 'Papua New Guinea',
+            'PH': 'Philippines',
+            'PK': 'Pakistan',
+            'PL': 'Poland',
+            'PR': 'Puerto Rico',
+            'PS': 'Palestine',
+            'PT': 'Portugal',
+            'PW': 'Palau',
+            'PY': 'Paraguay',
+            'QA': 'Qatar',
+            'RE': 'Reunion',
+            'RO': 'Romania',
+            'RS': 'Serbia',
+            'RU': 'Russia',
+            'RW': 'Rwanda',
+            'SA': 'Saudi Arabia',
+            'SB': 'Solomon Islands',
+            'SC': 'Seychelles',
+            'SD': 'Sudan',
+            'SE': 'Sweden',
+            'SG': 'Singapore',
+            'SH': 'St. Helena',
+            'SI': 'Slovenia',
+            'SJ': 'Svalbard and Jan Mayen Islands',
+            'SK': 'Slovakia',
+            'SL': 'Sierra Leone',
+            'SM': 'San Marino',
+            'SN': 'Senegal',
+            'SO': 'Somalia',
+            'SS': 'South Sudan',
+            'SV': 'El Salvador',
+            'SX': 'Sint Maarten',
+            'SY': 'Syria',
+            'SZ': 'Eswatini',
+            'TG': 'Togo',
+            'TH': 'Thailand',
+            'TJ': 'Tajikistan',
+            'TL': 'Timor-Leste',
+            'TM': 'Turkmenistan',
+            'TN': 'Tunisia',
+            'TO': 'Tonga',
+            'TR': 'Turkey',
+            'TT': 'Trinidad and Tobago',
+            'TW': 'Taiwan',
+            'TZ': 'Tanzania',
+            'UA': 'Ukraine',
+            'UG': 'Uganda',
+            'US': 'United States',
+            'UY': 'Uruguay',
+            'UZ': 'Uzbekistan',
+            'VA': 'Vatican',
+            'VE': 'Venezuela',
+            'VG': 'British Virgin Islands',
+            'VI': 'United States Virgin Islands',
+            'VN': 'Vietnam',
+            'VU': 'Vanuatu',
+            'WS': 'Samoa',
+            'XK': 'Kosovo',
+            'YE': 'Yemen',
+            'ZA': 'South Africa',
+            'ZM': 'Zambia',
+            'ZW': 'Zimbabwe',
+        }

MTIL_datasets/dtd.py

@@ -0,0 +1,95 @@
+import os
+import pickle
+import random
+
+from .utils import *
+
+from .oxford_pets import OxfordPets
+
+
+class DescribableTextures(DatasetBase):
+
+    dataset_dir = "dtd"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.image_dir = os.path.join(self.dataset_dir, "images")
+        self.split_path = os.path.join(self.dataset_dir, "split_zhou_DescribableTextures.json")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        if os.path.exists(self.split_path):
+            train, val, test = OxfordPets.read_split(self.split_path, self.image_dir)
+        else:
+            train, val, test = self.read_and_split_data(self.image_dir)
+            OxfordPets.save_split(train, val, test, self.split_path, self.image_dir)
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        subsample = subsample_classes
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample)
+
+        self.templates = [
+            lambda c: f'a photo of a {c} texture.',
+            lambda c: f'a photo of a {c} pattern.',
+            lambda c: f'a photo of a {c} thing.',
+            lambda c: f'a photo of a {c} object.',
+            lambda c: f'a photo of the {c} texture.',
+            lambda c: f'a photo of the {c} pattern.',
+            lambda c: f'a photo of the {c} thing.',
+            lambda c: f'a photo of the {c} object.',
+        ]
+        
+        super().__init__(train_x=train, val=val, test=test)
+
+    @staticmethod
+    def read_and_split_data(image_dir, p_trn=0.5, p_val=0.2, ignored=[], new_cnames=None):
+        categories = listdir_nohidden(image_dir)
+        categories = [c for c in categories if c not in ignored]
+        categories.sort()
+
+        p_tst = 1 - p_trn - p_val
+        print(f"Splitting into {p_trn:.0%} train, {p_val:.0%} val, and {p_tst:.0%} test")
+
+        def _collate(ims, y, c):
+            items = []
+            for im in ims:
+                item = Datum(impath=im, label=y, classname=c)
+                items.append(item)
+            return items
+
+        train, val, test = [], [], []
+        for label, category in enumerate(categories):
+            category_dir = os.path.join(image_dir, category)
+            images = listdir_nohidden(category_dir)
+            images = [os.path.join(category_dir, im) for im in images]
+            random.shuffle(images)
+            n_total = len(images)
+            n_train = round(n_total * p_trn)
+            n_val = round(n_total * p_val)
+            n_test = n_total - n_train - n_val
+            assert n_train > 0 and n_val > 0 and n_test > 0
+
+            if new_cnames is not None and category in new_cnames:
+                category = new_cnames[category]
+
+            train.extend(_collate(images[:n_train], label, category))
+            val.extend(_collate(images[n_train : n_train + n_val], label, category))
+            test.extend(_collate(images[n_train + n_val :], label, category))
+
+        return train, val, test

MTIL_datasets/eurosat.py

@@ -0,0 +1,75 @@
+import os
+import pickle
+
+from .utils import *
+
+from .oxford_pets import OxfordPets
+from .dtd import DescribableTextures as DTD
+
+NEW_CNAMES = {
+    "AnnualCrop": "Annual Crop Land",
+    "Forest": "Forest",
+    "HerbaceousVegetation": "Herbaceous Vegetation Land",
+    "Highway": "Highway or Road",
+    "Industrial": "Industrial Buildings",
+    "Pasture": "Pasture Land",
+    "PermanentCrop": "Permanent Crop Land",
+    "Residential": "Residential Buildings",
+    "River": "River",
+    "SeaLake": "Sea or Lake",
+}
+
+
+class EuroSAT(DatasetBase):
+
+    dataset_dir = "eurosat"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.image_dir = os.path.join(self.dataset_dir, "2750")
+        self.split_path = os.path.join(self.dataset_dir, "split_zhou_EuroSAT.json")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        if os.path.exists(self.split_path):
+            train, val, test = OxfordPets.read_split(self.split_path, self.image_dir)
+        else:
+            train, val, test = DTD.read_and_split_data(self.image_dir, new_cnames=NEW_CNAMES)
+            OxfordPets.save_split(train, val, test, self.split_path, self.image_dir)
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        subsample = subsample_classes
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample)
+
+        self.templates = [
+            lambda c: f"a centered satellite photo of {c}.",
+            lambda c: f"a centered satellite photo of a {c}.",
+            lambda c: f"a centered satellite photo of the {c}.",
+        ]
+
+        super().__init__(train_x=train, val=val, test=test)
+
+    def update_classname(self, dataset_old):
+        dataset_new = []
+        for item_old in dataset_old:
+            cname_old = item_old.classname
+            cname_new = NEW_CLASSNAMES[cname_old]
+            item_new = Datum(impath=item_old.impath, label=item_old.label, classname=cname_new)
+            dataset_new.append(item_new)
+        return dataset_new

MTIL_datasets/fer2013.py

@@ -0,0 +1,170 @@
+import os
+import pickle
+import random
+import warnings
+from collections import defaultdict
+
+from .utils import *
+from .oxford_pets import OxfordPets
+
+# Synonym sets per class (label index is the list index)
+FER2013_CLASSES_SYNONYMS = [
+    ['angry'],
+    ['disgusted', 'disgust'],
+    ['fearful', 'fear'],
+    ['happy', 'smiling'],
+    ['sad', 'depressed'],
+    ['surprised', 'surprise', 'shocked', 'spooked'],
+    ['neutral', 'bored'],
+]
+
+# Canonical class names are the first synonym in each list
+FER2013_CANONICAL = [syns[0] for syns in FER2013_CLASSES_SYNONYMS]
+
+# Prompt templates
+FER2013_TEMPLATES = [
+    'a photo of a {} looking face.',
+    'a photo of a face showing the emotion: {}.',
+    'a photo of a face looking {}.',
+    'a face that looks {}.',
+    'they look {}.',
+    'look at how {} they are.',
+]
+
+FER2013_DEBUG = os.environ.get("FER2013_DEBUG", "0") not in ("0", "false", "False", "")
+
+def _dbg(msg: str):
+    if FER2013_DEBUG:
+        print(f"[FER2013][DEBUG] {msg}")
+
+
+def _norm(s: str) -> str:
+    s = s.lower().strip()
+    for ch in [" ", "_", "-", "."]:
+        s = s.replace(ch, "")
+    return s
+
+
+def _build_syn_map():
+    m = {}
+    for y, syns in enumerate(FER2013_CLASSES_SYNONYMS):
+        for s in syns:
+            m[_norm(s)] = y
+    # add common canonical variants for safety
+    aliases = {
+        'disgust': 1,
+        'fear': 2,
+        'surprise': 5,
+    }
+    for k, v in aliases.items():
+        m[_norm(k)] = v
+    return m
+
+
+class FER2013(DatasetBase):
+
+    dataset_dir = "fer2013"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.split_path = os.path.join(self.dataset_dir, "split_custom_FER2013.json")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        train_dir = os.path.join(self.dataset_dir, "train")
+        test_dir = os.path.join(self.dataset_dir, "test")
+        if not os.path.isdir(train_dir) or not os.path.isdir(test_dir):
+            raise ValueError(
+                f"FER2013: expected train/test folders under '{self.dataset_dir}'. Got train={os.path.isdir(train_dir)}, test={os.path.isdir(test_dir)}"
+            )
+
+        # try cache
+        if os.path.exists(self.split_path):
+            try:
+                train, val, test = OxfordPets.read_split(self.split_path, self.dataset_dir)
+            except Exception as e:
+                warnings.warn(f"FER2013: failed to read cached split; rebuilding. Error: {e}")
+                train, val, test = self._build_split(train_dir, test_dir)
+                try:
+                    OxfordPets.save_split(train, val, test, self.split_path, self.dataset_dir)
+                except Exception as e2:
+                    warnings.warn(f"FER2013: failed to save split: {e2}")
+        else:
+            train, val, test = self._build_split(train_dir, test_dir)
+            try:
+                OxfordPets.save_split(train, val, test, self.split_path, self.dataset_dir)
+            except Exception as e:
+                warnings.warn(f"FER2013: failed to save split: {e}")
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample_classes)
+        self.templates = FER2013_TEMPLATES
+        super().__init__(train_x=train, val=val, test=test)
+
+    def _build_split(self, train_dir, test_dir, p_val=0.2):
+        syn_map = _build_syn_map()
+        # read train per class
+        tr_items_by_label = defaultdict(list)
+        class_dirs = listdir_nohidden(train_dir, sort=True)
+        if not class_dirs:
+            warnings.warn(f"FER2013: no class folders found in {train_dir}")
+        for cls in class_dirs:
+            full = os.path.join(train_dir, cls)
+            if not os.path.isdir(full):
+                continue
+            key = _norm(cls)
+            y = syn_map.get(key)
+            if y is None:
+                warnings.warn(f"FER2013: unexpected class folder in train: '{cls}'")
+                continue
+            cname = FER2013_CANONICAL[y]
+            for fname in listdir_nohidden(full):
+                impath = os.path.join(full, fname)
+                tr_items_by_label[y].append(Datum(impath=impath, label=y, classname=cname))
+
+        # stratified split train->(train,val)
+        train, val = [], []
+        for y, items in tr_items_by_label.items():
+            random.shuffle(items)
+            n_val = max(1, round(len(items) * p_val)) if len(items) > 1 else 0
+            val.extend(items[:n_val])
+            train.extend(items[n_val:])
+
+        # read test
+        test = []
+        class_dirs = listdir_nohidden(test_dir, sort=True)
+        if not class_dirs:
+            warnings.warn(f"FER2013: no class folders found in {test_dir}")
+        for cls in class_dirs:
+            full = os.path.join(test_dir, cls)
+            if not os.path.isdir(full):
+                continue
+            key = _norm(cls)
+            y = syn_map.get(key)
+            if y is None:
+                warnings.warn(f"FER2013: unexpected class folder in test: '{cls}'")
+                continue
+            cname = FER2013_CANONICAL[y]
+            for fname in listdir_nohidden(full):
+                impath = os.path.join(full, fname)
+                test.append(Datum(impath=impath, label=y, classname=cname))
+
+        # basic sanity
+        if not train or not val or not test:
+            warnings.warn(f"FER2013: split sizes train={len(train)} val={len(val)} test={len(test)}")
+        return train, val, test

MTIL_datasets/fgvc_aircraft.py

@@ -0,0 +1,73 @@
+import os
+import pickle
+
+from .utils import *
+
+from .oxford_pets import OxfordPets
+import torchvision.datasets
+
+
+class FGVCAircraft(DatasetBase):
+
+    dataset_dir = "fgvc_aircraft"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.image_dir = os.path.join(self.dataset_dir, "images")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        classnames = []
+        with open(os.path.join(self.dataset_dir, "variants.txt"), "r") as f:
+            lines = f.readlines()
+            for line in lines:
+                classnames.append(line.strip())
+        cname2lab = {c: i for i, c in enumerate(classnames)}
+
+        train = self.read_data(cname2lab, "images_variant_train.txt")
+        val = self.read_data(cname2lab, "images_variant_val.txt")
+        test = self.read_data(cname2lab, "images_variant_test.txt")
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        subsample = subsample_classes
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample)
+
+        self.templates = [
+            lambda c: f"a photo of a {c}, a type of aircraft.",
+            lambda c: f"a photo of the {c}, a type of aircraft.",
+        ]
+
+        super().__init__(train_x=train, val=val, test=test)
+
+    def read_data(self, cname2lab, split_file):
+        filepath = os.path.join(self.dataset_dir, split_file)
+        items = []
+
+        with open(filepath, "r") as f:
+            lines = f.readlines()
+            for line in lines:
+                line = line.strip().split(" ")
+                imname = line[0] + ".jpg"
+                classname = " ".join(line[1:])
+                impath = os.path.join(self.image_dir, imname)
+                label = cname2lab[classname]
+                item = Datum(impath=impath, label=label, classname=classname)
+                items.append(item)
+
+        return items

MTIL_datasets/food101.py

@@ -0,0 +1,52 @@
+import os
+import pickle
+
+from .utils import *
+
+from .oxford_pets import OxfordPets
+from .dtd import DescribableTextures as DTD
+
+
+
+class Food101(DatasetBase):
+
+    dataset_dir = "food-101"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.image_dir = os.path.join(self.dataset_dir, "images")
+        self.split_path = os.path.join(self.dataset_dir, "split_zhou_Food101.json")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        if os.path.exists(self.split_path):
+            train, val, test = OxfordPets.read_split(self.split_path, self.image_dir)
+        else:
+            train, val, test = DTD.read_and_split_data(self.image_dir)
+            OxfordPets.save_split(train, val, test, self.split_path, self.image_dir)
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        subsample = subsample_classes
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample)
+
+        self.templates = [
+            lambda c: f"a photo of a {c}, a type of food.",
+        ]
+
+        super().__init__(train_x=train, val=val, test=test)

MTIL_datasets/gtsrb.py

@@ -0,0 +1,202 @@
+import os
+import pickle
+import random
+import warnings
+
+from .utils import *
+from .oxford_pets import OxfordPets
+
+# 43 classes for GTSRB with human-readable names
+classes = [
+    'red and white circle 20 kph speed limit',
+    'red and white circle 30 kph speed limit',
+    'red and white circle 50 kph speed limit',
+    'red and white circle 60 kph speed limit',
+    'red and white circle 70 kph speed limit',
+    'red and white circle 80 kph speed limit',
+    'end / de-restriction of 80 kph speed limit',
+    'red and white circle 100 kph speed limit',
+    'red and white circle 120 kph speed limit',
+    'red and white circle red car and black car no passing',
+    'red and white circle red truck and black car no passing',
+    'red and white triangle road intersection warning',
+    'white and yellow diamond priority road',
+    'red and white upside down triangle yield right-of-way',
+    'stop',
+    'empty red and white circle',
+    'red and white circle no truck entry',
+    'red circle with white horizonal stripe no entry',
+    'red and white triangle with exclamation mark warning',
+    'red and white triangle with black left curve approaching warning',
+    'red and white triangle with black right curve approaching warning',
+    'red and white triangle with black double curve approaching warning',
+    'red and white triangle rough / bumpy road warning',
+    'red and white triangle car skidding / slipping warning',
+    'red and white triangle with merging / narrow lanes warning',
+    'red and white triangle with person digging / construction / road work warning',
+    'red and white triangle with traffic light approaching warning',
+    'red and white triangle with person walking warning',
+    'red and white triangle with child and person walking warning',
+    'red and white triangle with bicyle warning',
+    'red and white triangle with snowflake / ice warning',
+    'red and white triangle with deer warning',
+    'white circle with gray strike bar no speed limit',
+    'blue circle with white right turn arrow mandatory',
+    'blue circle with white left turn arrow mandatory',
+    'blue circle with white forward arrow mandatory',
+    'blue circle with white forward or right turn arrow mandatory',
+    'blue circle with white forward or left turn arrow mandatory',
+    'blue circle with white keep right arrow mandatory',
+    'blue circle with white keep left arrow mandatory',
+    'blue circle with white arrows indicating a traffic circle',
+    'white circle with gray strike bar indicating no passing for cars has ended',
+    'white circle with gray strike bar indicating no passing for trucks has ended',
+]
+
+
+DEBUG = os.environ.get("GTSRB_DEBUG", "0") not in ("0", "false", "False", "")
+
+
+def _dbg(msg: str):
+    if DEBUG:
+        print(f"[GTSRB][DEBUG] {msg}")
+
+
+def _is_image_file(name):
+    name = name.lower()
+    return any(name.endswith(ext) for ext in ['.ppm', '.png', '.jpg', '.jpeg', '.bmp', '.webp'])
+
+
+class GTSRB(DatasetBase):
+    """
+    German Traffic Sign Recognition Benchmark (GTSRB)
+
+    Expected structure:
+      <root>/gtsrb/
+        ├─ 00000/*.ppm
+        ├─ 00001/*.ppm
+        └─ ... up to 00042/
+
+    Note: The dataset has no official test split in this layout, so we will
+    randomly split each class into train/val/test (50%/20%/30%), cached to json.
+    """
+
+    dataset_dir = "gtsrb"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.image_dir = self.dataset_dir
+        self.split_path = os.path.join(self.dataset_dir, "split_custom_GTSRB.json")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+        _dbg(f"dataset_dir={self.dataset_dir}")
+        _dbg(f"image_dir={self.image_dir}")
+        _dbg(f"split_path exists? {os.path.exists(self.split_path)}")
+
+        if os.path.exists(self.split_path):
+            try:
+                train, val, test = OxfordPets.read_split(self.split_path, self.image_dir)
+            except Exception as e:
+                warnings.warn(f"GTSRB: failed to read split file '{self.split_path}'; rebuilding split. Error: {e}")
+                train, val, test = self.read_and_split_data(self.image_dir)
+                try:
+                    OxfordPets.save_split(train, val, test, self.split_path, self.image_dir)
+                except Exception as e2:
+                    warnings.warn(f"GTSRB: failed to save rebuilt split to '{self.split_path}': {e2}")
+        else:
+            train, val, test = self.read_and_split_data(self.image_dir)
+            try:
+                OxfordPets.save_split(train, val, test, self.split_path, self.image_dir)
+            except Exception as e:
+                warnings.warn(f"GTSRB: failed to save split to '{self.split_path}': {e}")
+
+        _dbg(f"loaded counts: train={len(train)}, val={len(val)}, test={len(test)}")
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+            _dbg(f"few-shot applied: train={len(train)}, val={len(val)} (shots={num_shots})")
+
+        # Ensure class subsampling behavior matches others
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample_classes)
+        _dbg(f"after subsample='{subsample_classes}': train={len(train)}, val={len(val)}, test={len(test)})")
+
+        # Prompt templates provided by user
+        self.templates = [
+            lambda c: f'a zoomed in photo of a "{c}" traffic sign.',
+            lambda c: f'a centered photo of a "{c}" traffic sign.',
+            lambda c: f'a close up photo of a "{c}" traffic sign.',
+        ]
+
+        super().__init__(train_x=train, val=val, test=test)
+
+    @staticmethod
+    def read_and_split_data(image_root, p_trn=0.5, p_val=0.2):
+        # Discover class directories (e.g., 00000 .. 00042)
+        try:
+            all_entries = listdir_nohidden(image_root, sort=True)
+        except Exception as e:
+            warnings.warn(f"GTSRB: failed to list directory '{image_root}': {e}")
+            raise
+        class_dirs = [d for d in all_entries if os.path.isdir(os.path.join(image_root, d))]
+        class_dirs.sort()
+        assert len(class_dirs) > 0, f"GTSRB: no class folders found under {image_root}"
+        _dbg(f"found {len(class_dirs)} class folders; head={class_dirs[:5]}")
+
+        # Map sorted class dirs to labels 0..N-1 and names from 'classes'
+        if len(class_dirs) != len(classes):
+            print(f"Warning: detected {len(class_dirs)} class folders but classes list has {len(classes)} entries. Proceeding with min overlap.")
+        num_labels = min(len(class_dirs), len(classes))
+
+        def _collate(paths, y, cname):
+            return [Datum(impath=p, label=y, classname=cname) for p in paths]
+
+        train, val, test = [], [], []
+        for y, cls_dir in enumerate(class_dirs[:num_labels]):
+            cname = classes[y]
+            cdir = os.path.join(image_root, cls_dir)
+            try:
+                files = listdir_nohidden(cdir, sort=False)
+            except Exception as e:
+                warnings.warn(f"GTSRB: failed to list class folder '{cdir}': {e}")
+                files = []
+            imgs = [os.path.join(cdir, f) for f in files if _is_image_file(f)]
+            random.shuffle(imgs)
+            n_total = len(imgs)
+            if n_total == 0:
+                warnings.warn(f"GTSRB: empty or unreadable class folder: {cdir}; skipping this class")
+                continue
+            if n_total < 5:
+                warnings.warn(f"GTSRB: very few images in class '{cls_dir}' (n={n_total}); splits may be unstable")
+            n_train = round(n_total * p_trn)
+            n_val = round(n_total * p_val)
+            n_test = n_total - n_train - n_val
+            if not (n_train > 0 and n_val > 0 and n_test > 0):
+                warnings.warn(f"GTSRB: split would create empty split for class {cls_dir} (n={n_total}); adjusting strategy to keep at least 1 per split")
+                # Fallback: enforce at least 1 per split if possible
+                if n_total >= 3:
+                    n_train, n_val, n_test = 1, 1, n_total - 2
+                elif n_total == 2:
+                    n_train, n_val, n_test = 1, 1, 0
+                else:  # n_total == 1
+                    n_train, n_val, n_test = 1, 0, 0
+
+            train.extend(_collate(imgs[:n_train], y, cname))
+            val.extend(_collate(imgs[n_train:n_train + n_val], y, cname))
+            test.extend(_collate(imgs[n_train + n_val:], y, cname))
+            _dbg(f"class {cls_dir} -> label {y}: total={n_total}, train={n_train}, val={n_val}, test={n_test}")
+
+        _dbg(f"aggregate sizes: train={len(train)}, val={len(val)}, test={len(test)}")
+        return train, val, test

MTIL_datasets/hatefulmemes.py

@@ -0,0 +1,117 @@
+import os
+import json
+import pickle
+
+from .utils import *
+from .oxford_pets import OxfordPets
+
+
+class HatefulMemes(DatasetBase):
+    """
+    Hateful Memes dataset loader (image-only for MTIL).
+
+    Expected structure:
+      <root>/hatefulmemes/
+        ├─ img/*.png|jpg
+        ├─ train.jsonl
+        ├─ dev.jsonl
+        └─ test.jsonl
+
+    JSONL lines example:
+      {"id":85362, "img":"img/85362.png", "label":0, "text":"..."}
+
+    We DO NOT use the 'text' field. Only image path and label are used.
+
+    Classes: ['meme', 'hatespeech meme']
+    Templates: ['a {}.']
+    """
+
+    dataset_dir = "hatefulmemes"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        if not os.path.isdir(self.dataset_dir):
+            # allow alias without trailing 's'
+            alt = os.path.join(root, "hatefulmeme")
+            if os.path.isdir(alt):
+                self.dataset_dir = alt
+            else:
+                raise ValueError(
+                    f"HatefulMemes: dataset folder not found at '{self.dataset_dir}' or '{alt}'"
+                )
+
+        split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(split_fewshot_dir)
+
+        # paths
+        train_json = os.path.join(self.dataset_dir, "train.jsonl")
+        dev_json = os.path.join(self.dataset_dir, "dev.jsonl")
+        test_json = os.path.join(self.dataset_dir, "test.jsonl")
+
+        if not os.path.isfile(train_json) or not os.path.isfile(dev_json) or not os.path.isfile(test_json):
+            raise ValueError(
+                f"HatefulMemes: missing jsonl files. Expected train/dev/test at '{self.dataset_dir}'."
+            )
+
+        # fixed classes
+        classes = ["meme", "hatespeech meme"]
+        lab_to_name = {0: classes[0], 1: classes[1]}
+
+        # read splits
+        train = self._read_jsonl(train_json, lab_to_name)
+        val = self._read_jsonl(dev_json, lab_to_name)
+        test = self._read_jsonl(test_json, lab_to_name)
+        # Some public releases of Hateful Memes do not include labels for test.jsonl.
+        # In that case, fallback to use dev.jsonl for evaluation so zero-shot works.
+        if len(test) == 0 and len(val) > 0:
+            test = list(val)
+
+        # few-shot
+        if num_shots >= 1:
+            preprocessed = os.path.join(split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            if os.path.exists(preprocessed):
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        # subsample behavior consistent with others
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample_classes)
+
+        # templates per user spec
+        self.templates = [
+            lambda c: f'a {c}.',
+        ]
+
+        super().__init__(train_x=train, val=val, test=test)
+
+    def _read_jsonl(self, filepath, lab_to_name):
+        items = []
+        with open(filepath, 'r', encoding='utf-8') as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    obj = json.loads(line)
+                except Exception:
+                    continue
+                # fields: id, img, label, text (ignored)
+                img_rel = obj.get('img', '')
+                label = obj.get('label', None)
+                if img_rel is None or label is None:
+                    continue
+                if label not in (0, 1):
+                    raise ValueError(f"HatefulMemes: unexpected label {label} in {filepath}")
+                impath = os.path.join(self.dataset_dir, img_rel.replace('/', os.sep))
+                if not os.path.isfile(impath):
+                    continue
+                cname = lab_to_name[label]
+                items.append(Datum(impath=impath, label=label, classname=cname))
+        return items

MTIL_datasets/imagenet_r.py

@@ -0,0 +1,231 @@
+import os
+import random
+from typing import List, Tuple, Dict
+
+from .utils import *  # Datum, DatasetBase, listdir_nohidden, mkdir_if_missing
+from .oxford_pets import OxfordPets
+
+
+IMAGENETR_TEMPLATES: List[str] = [
+    'a bad photo of a {}.',
+    'a photo of many {}.',
+    'a sculpture of a {}.',
+    'a photo of the hard to see {}.',
+    'a low resolution photo of the {}.',
+    'a rendering of a {}.',
+    'graffiti of a {}.',
+    'a bad photo of the {}.',
+    'a cropped photo of the {}.',
+    'a tattoo of a {}.',
+    'the embroidered {}.',
+    'a photo of a hard to see {}.',
+    'a bright photo of a {}.',
+    'a photo of a clean {}.',
+    'a photo of a dirty {}.',
+    'a dark photo of the {}.',
+    'a drawing of a {}.',
+    'a photo of my {}.',
+    'the plastic {}.',
+    'a photo of the cool {}.',
+    'a close-up photo of a {}.',
+    'a black and white photo of the {}.',
+    'a painting of the {}.',
+    'a painting of a {}.',
+    'a pixelated photo of the {}.',
+    'a sculpture of the {}.',
+    'a bright photo of the {}.',
+    'a cropped photo of a {}.',
+    'a plastic {}.',
+    'a photo of the dirty {}.',
+    'a jpeg corrupted photo of a {}.',
+    'a blurry photo of the {}.',
+    'a photo of the {}.',
+    'a good photo of the {}.',
+    'a rendering of the {}.',
+    'a {} in a video game.',
+    'a photo of one {}.',
+    'a doodle of a {}.',
+    'a close-up photo of the {}.',
+    'a photo of a {}.',
+    'the origami {}.',
+    'the {} in a video game.',
+    'a sketch of a {}.',
+    'a doodle of the {}.',
+    'a origami {}.',
+    'a low resolution photo of a {}.',
+    'the toy {}.',
+    'a rendition of the {}.',
+    'a photo of the clean {}.',
+    'a photo of a large {}.',
+    'a rendition of a {}.',
+    'a photo of a nice {}.',
+    'a photo of a weird {}.',
+    'a blurry photo of a {}.',
+    'a cartoon {}.',
+    'art of a {}.',
+    'a sketch of the {}.',
+    'a embroidered {}.',
+    'a pixelated photo of a {}.',
+    'itap of the {}.',
+    'a jpeg corrupted photo of the {}.',
+    'a good photo of a {}.',
+    'a plushie {}.',
+    'a photo of the nice {}.',
+    'a photo of the small {}.',
+    'a photo of the weird {}.',
+    'the cartoon {}.',
+    'art of the {}.',
+    'a drawing of the {}.',
+    'a photo of the large {}.',
+    'a black and white photo of a {}.',
+    'the plushie {}.',
+    'a dark photo of a {}.',
+    'itap of a {}.',
+    'graffiti of the {}.',
+    'a toy {}.',
+    'itap of my {}.',
+    'a photo of a cool {}.',
+    'a photo of a small {}.',
+    'a tattoo of the {}.',
+]
+
+
+class ImageNetR(DatasetBase):
+
+    dataset_dir = "imagenet-r"
+
+    def __init__(self, root: str, num_shots: int = 0, seed: int = 1, subsample_classes: str = 'all', test_ratio: float = 0.2):
+        """
+        Expect directory structure:
+            {root}/imagenet-r/<wnid>/*.jpg|png|jpeg|bmp|webp
+            {root}/imagenet-r/classname.txt  # lines: "<wnid> <human_readable_name>"
+        There is no official test split; we perform a per-class split into train/test, then
+        split train into train/val using OxfordPets.split_trainval.
+        """
+        rnd = random.Random(seed)
+
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.split_path = os.path.join(self.dataset_dir, "split_custom_ImageNetR.json")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        # Read class mapping from classname.txt
+        class_map_path = os.path.join(self.dataset_dir, "classname.txt")
+        if not os.path.exists(class_map_path):
+            raise FileNotFoundError(f"Class mapping file not found: {class_map_path}")
+        wnids, classnames = self._read_class_map(class_map_path)
+        self._wnids = wnids
+        self._classnames_ref = classnames
+
+        # Build or load split
+        if os.path.exists(self.split_path):
+            train, val, test = OxfordPets.read_split(self.split_path, self.dataset_dir)
+        else:
+            trainval, test = self._read_from_folders(self.dataset_dir, wnids, classnames, rnd=rnd, test_ratio=test_ratio)
+            train, val = self._split_trainval_safe(trainval)
+            OxfordPets.save_split(train, val, test, self.split_path, self.dataset_dir)
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                import pickle
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                import pickle
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample_classes)
+
+        # Expose templates
+        self.templates = IMAGENETR_TEMPLATES
+
+        super().__init__(train_x=train, val=val, test=test)
+
+    @staticmethod
+    def _read_class_map(filepath: str) -> Tuple[List[str], List[str]]:
+        wnids: List[str] = []
+        cnames: List[str] = []
+        with open(filepath, 'r', encoding='utf-8') as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                parts = line.split()
+                wnid = parts[0]
+                cname = ' '.join(parts[1:]) if len(parts) > 1 else wnid
+                wnids.append(wnid)
+                cnames.append(cname)
+        return wnids, cnames
+
+    @staticmethod
+    def _split_trainval_safe(trainval: List[Datum], p_val: float = 0.2) -> Tuple[List[Datum], List[Datum]]:
+        from collections import defaultdict
+        tracker: Dict[int, List[int]] = defaultdict(list)
+        for idx, item in enumerate(trainval):
+            tracker[item.label].append(idx)
+
+        train, val = [], []
+        for label, idxs in tracker.items():
+            n = len(idxs)
+            if n <= 1:
+                for idx in idxs:
+                    train.append(trainval[idx])
+                continue
+            n_val = max(1, int(round(n * p_val)))
+            if n_val >= n:
+                n_val = n - 1
+            random.shuffle(idxs)
+            for i, idx in enumerate(idxs):
+                if i < n_val:
+                    val.append(trainval[idx])
+                else:
+                    train.append(trainval[idx])
+        if len(val) == 0 and len(train) > 0:
+            val.append(train[-1])
+            train = train[:-1]
+        return train, val
+
+    def _read_from_folders(self, data_dir: str, wnids: List[str], classnames: List[str], rnd: random.Random, test_ratio: float) -> Tuple[List[Datum], List[Datum]]:
+        exts = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}
+        class_to_label = {wnid: i for i, wnid in enumerate(wnids)}
+        label_to_cname = {i: classnames[i] for i in range(len(classnames))}
+
+        items_by_label: Dict[int, List[Datum]] = {i: [] for i in range(len(wnids))}
+        for wnid in wnids:
+            cdir = os.path.join(data_dir, wnid)
+            if not os.path.isdir(cdir):
+                # If a class listed in mapping has no folder, skip gracefully
+                continue
+            label = class_to_label[wnid]
+            cname = label_to_cname[label]
+            for fname in listdir_nohidden(cdir, sort=True):
+                fext = os.path.splitext(fname)[1].lower()
+                if fext not in exts:
+                    continue
+                impath = os.path.join(cdir, fname)
+                items_by_label[label].append(Datum(impath=impath, label=label, classname=cname))
+
+        trainval: List[Datum] = []
+        test: List[Datum] = []
+        for label, items in items_by_label.items():
+            if not items:
+                continue
+            rnd.shuffle(items)
+            if len(items) == 1:
+                trainval.extend(items)
+                continue
+            n_test = max(1, int(round(len(items) * test_ratio)))
+            if n_test >= len(items):
+                n_test = len(items) - 1
+            test.extend(items[:n_test])
+            trainval.extend(items[n_test:])
+        return trainval, test

MTIL_datasets/kitti_distance.py

@@ -0,0 +1,228 @@
+import os
+import random
+from typing import List, Tuple
+
+from .utils import *  # Datum, DatasetBase, listdir_nohidden, mkdir_if_missing
+from .oxford_pets import OxfordPets
+
+# Class names per user specification (ordered)
+KITTI_DISTANCE_CLASSES: List[str] = [
+    'a photo i took of a car nearby',
+    'a photo i took with a car in the middle distance',
+    'a photo i took with a car faraway',
+    'a photo i took with no car.',
+]
+
+# Keep templates as strings with {}
+KITTI_DISTANCE_TEMPLATES: List[str] = [
+    '{}',
+]
+
+
+class KittiDistance(DatasetBase):
+
+    dataset_dir = "kitti"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all', test_ratio: float = 0.2):
+        """
+        Predict the distance category of the closest car from KITTI labels.
+
+        Expected structure:
+            {root}/kitti/
+                image/000004.png
+                label/000004.txt
+
+        Label file format: KITTI object detection label lines.
+        We retain only lines with type == 'Car'.
+        For these, we read the 14th column (index 13, loc_z) as the camera-depth in meters.
+        We take the minimum positive loc_z across all 'Car' lines as z_min for the image.
+        If there is no valid positive loc_z for 'Car', we assign the 'no car' class.
+
+        Discretization into classes:
+            0: 0 < z_min < 10         -> 'nearby'
+            1: 10 <= z_min < 30       -> 'middle distance'
+            2: z_min >= 30            -> 'faraway'
+            3: no car                 -> 'no car.'
+
+        No official test set; we split randomly (per class) into train/test by test_ratio.
+        Then we split train into train/val using a safe per-class split.
+        Splits are saved to JSON for reproducibility.
+        """
+        rnd = random.Random(seed)
+
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+
+        # Image and label directories
+        image_dir = os.path.join(self.dataset_dir, 'image')
+        if not os.path.isdir(image_dir):
+            # Graceful fallback to standard KITTI naming variants if provided
+            alt1 = os.path.join(self.dataset_dir, 'image_2')
+            alt2 = os.path.join(self.dataset_dir, 'images')
+            if os.path.isdir(alt1):
+                image_dir = alt1
+            elif os.path.isdir(alt2):
+                image_dir = alt2
+        label_dir = os.path.join(self.dataset_dir, 'label')
+        if not os.path.isdir(label_dir):
+            alt_l1 = os.path.join(self.dataset_dir, 'label_2')
+            alt_l2 = os.path.join(self.dataset_dir, 'labels')
+            if os.path.isdir(alt_l1):
+                label_dir = alt_l1
+            elif os.path.isdir(alt_l2):
+                label_dir = alt_l2
+
+        if not os.path.isdir(image_dir):
+            raise FileNotFoundError(f"KittiDistance: image dir not found: {image_dir}")
+        if not os.path.isdir(label_dir):
+            raise FileNotFoundError(f"KittiDistance: label dir not found: {label_dir}")
+
+        self.image_dir = image_dir
+        self.label_dir = label_dir
+
+        self.split_path = os.path.join(self.dataset_dir, "split_custom_KittiDistance.json")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        if os.path.exists(self.split_path):
+            train, val, test = OxfordPets.read_split(self.split_path, self.image_dir)
+        else:
+            trainval, test = self._read_and_split(self.image_dir, self.label_dir, rnd=rnd, test_ratio=test_ratio)
+            train, val = self._split_trainval_safe(trainval)
+            OxfordPets.save_split(train, val, test, self.split_path, self.image_dir)
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                import pickle
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                import pickle
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        # Allow class subsampling if requested
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample_classes)
+
+        # Templates per user specification
+        self.templates = KITTI_DISTANCE_TEMPLATES
+
+        super().__init__(train_x=train, val=val, test=test)
+        # Override inferred metadata to ensure stable 4-way classification regardless of train label coverage
+        self._classnames = KITTI_DISTANCE_CLASSES
+        self._lab2cname = {i: c for i, c in enumerate(KITTI_DISTANCE_CLASSES)}
+        self._num_classes = len(KITTI_DISTANCE_CLASSES)
+
+    def _read_and_split(self, image_dir: str, label_dir: str, rnd: random.Random, test_ratio: float) -> Tuple[List[Datum], List[Datum]]:
+        exts = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}
+        # Collect items with computed labels
+        items_by_label = {i: [] for i in range(len(KITTI_DISTANCE_CLASSES))}
+
+        for fname in listdir_nohidden(image_dir, sort=True):
+            fext = os.path.splitext(fname)[1].lower()
+            if fext not in exts:
+                continue
+            impath = os.path.join(image_dir, fname)
+            stem = os.path.splitext(fname)[0]
+            label_path = os.path.join(label_dir, stem + '.txt')
+            label_idx = self._compute_label_from_kitti(label_path)
+            cname = KITTI_DISTANCE_CLASSES[label_idx]
+            items_by_label[label_idx].append(Datum(impath=impath, label=label_idx, classname=cname))
+
+        # Per-class balanced split into trainval/test
+        trainval, test = [], []
+        for label, items in items_by_label.items():
+            if not items:
+                continue
+            rnd.shuffle(items)
+            if len(items) == 1:
+                trainval.extend(items)
+                continue
+            n_test = max(1, int(round(len(items) * test_ratio)))
+            if n_test >= len(items):
+                n_test = len(items) - 1
+            test.extend(items[:n_test])
+            trainval.extend(items[n_test:])
+
+        return trainval, test
+
+    @staticmethod
+    def _compute_label_from_kitti(label_path: str) -> int:
+        """
+        Parse KITTI label file; focus only on lines with type == 'Car'.
+        Extract the 14th column (index 13, loc_z) as meters; consider only positive values.
+        If no positive loc_z found -> class 'no car' (index 3).
+
+        Thresholds:
+            0: 0 < z < 10
+            1: 10 <= z < 30
+            2: z >= 30
+        """
+        z_vals = []
+        if os.path.isfile(label_path):
+            with open(label_path, 'r') as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    parts = line.split()
+                    obj_type = parts[0]
+                    if obj_type != 'Car':
+                        continue
+                    if len(parts) < 14:
+                        # Not a valid KITTI detection line; skip
+                        continue
+                    try:
+                        # parts[13] is loc_z (0-based index), per KITTI format
+                        z = float(parts[13])
+                        if z > 0:
+                            z_vals.append(z)
+                    except Exception:
+                        continue
+        # Determine class index
+        if not z_vals:
+            return 3  # no car
+        z_min = min(z_vals)
+        if z_min < 10:
+            return 0
+        elif z_min < 30:
+            return 1
+        else:
+            return 2
+
+    @staticmethod
+    def _split_trainval_safe(trainval: List[Datum], p_val: float = 0.2) -> Tuple[List[Datum], List[Datum]]:
+        from collections import defaultdict
+        tracker = defaultdict(list)
+        for idx, item in enumerate(trainval):
+            tracker[item.label].append(idx)
+
+        train, val = [], []
+        for label, idxs in tracker.items():
+            n = len(idxs)
+            if n <= 1:
+                for idx in idxs:
+                    train.append(trainval[idx])
+                continue
+            n_val = max(1, int(round(n * p_val)))
+            if n_val >= n:
+                n_val = n - 1
+            random.shuffle(idxs)
+            for i, idx in enumerate(idxs):
+                if i < n_val:
+                    val.append(trainval[idx])
+                else:
+                    train.append(trainval[idx])
+
+        if len(val) == 0 and len(train) > 0:
+            val.append(train[-1])
+            train = train[:-1]
+
+        return train, val

MTIL_datasets/mnist.py

@@ -0,0 +1,132 @@
+import os
+import pickle
+
+from .utils import *
+
+from .oxford_pets import OxfordPets
+from .dtd import DescribableTextures as DTD
+from .oxford_pets import OxfordPets
+
+import torch
+import codecs
+import numpy as np
+import sys
+
+classes = [
+        "0 - zero",
+        "1 - one",
+        "2 - two",
+        "3 - three",
+        "4 - four",
+        "5 - five",
+        "6 - six",
+        "7 - seven",
+        "8 - eight",
+        "9 - nine",
+    ]
+
+class MNIST(DatasetBase):
+
+    dataset_dir = "mnist"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.image_dir = self.dataset_dir
+
+        trainval_image_file = "train-images-idx3-ubyte"
+        trainval_data = read_image_file(os.path.join(self.image_dir, trainval_image_file))  # Size([60000, 28, 28]) torch.uint8
+        trainval_label_file = "train-labels-idx1-ubyte"
+        trainval_targets = read_label_file(os.path.join(self.image_dir, trainval_label_file))  # Size([60000]) torch.int64
+        # trainval_names = 
+
+        test_image_file = "t10k-images-idx3-ubyte"
+        test_data = read_image_file(os.path.join(self.image_dir, test_image_file))
+        test_label_file = "t10k-labels-idx1-ubyte"
+        test_targets = read_label_file(os.path.join(self.image_dir, test_label_file))
+
+        trainval = []
+        for idx in range(trainval_data.size(0)):
+            item = Datum(impath=Image.fromarray(trainval_data[idx].numpy(), mode="L"), 
+                         label=int(trainval_targets[idx]), classname=classes[trainval_targets[idx]])
+            trainval.append(item)
+        
+        test = []
+        for idx in range(test_data.size(0)):
+            item = Datum(impath=Image.fromarray(test_data[idx].numpy(), mode="L"), 
+                         label=int(test_targets[idx]), classname=classes[test_targets[idx]])
+            test.append(item)
+
+        train, val = OxfordPets.split_trainval(trainval)
+
+        if num_shots >= 1:
+            train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+            val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+
+        subsample = subsample_classes
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample)
+
+        self.templates = [
+            lambda c: f'a photo of the number: "{c}".',
+        ]
+
+        super().__init__(train_x=train, val=val, test=test)
+
+
+def _flip_byte_order(t: torch.Tensor) -> torch.Tensor:
+    return (
+        t.contiguous().view(torch.uint8).view(*t.shape, t.element_size()).flip(-1).view(*t.shape[:-1], -1).view(t.dtype)
+    )
+
+def get_int(b: bytes) -> int:
+    return int(codecs.encode(b, "hex"), 16)
+
+SN3_PASCALVINCENT_TYPEMAP = {
+    8: torch.uint8,
+    9: torch.int8,
+    11: torch.int16,
+    12: torch.int32,
+    13: torch.float32,
+    14: torch.float64,
+}
+
+def read_sn3_pascalvincent_tensor(path: str, strict: bool = True) -> torch.Tensor:
+    """Read a SN3 file in "Pascal Vincent" format (Lush file 'libidx/idx-io.lsh').
+    Argument may be a filename, compressed filename, or file object.
+    """
+    # read
+    with open(path, "rb") as f:
+        data = f.read()
+    # parse
+    magic = get_int(data[0:4])
+    nd = magic % 256
+    ty = magic // 256
+    assert 1 <= nd <= 3
+    assert 8 <= ty <= 14
+    torch_type = SN3_PASCALVINCENT_TYPEMAP[ty]
+    s = [get_int(data[4 * (i + 1) : 4 * (i + 2)]) for i in range(nd)]
+
+    parsed = torch.frombuffer(bytearray(data), dtype=torch_type, offset=(4 * (nd + 1)))
+    if sys.byteorder == "little" and parsed.element_size() > 1:
+        parsed = _flip_byte_order(parsed)
+
+    assert parsed.shape[0] == np.prod(s) or not strict
+    return parsed.view(*s)
+
+
+def read_label_file(path: str) -> torch.Tensor:
+    x = read_sn3_pascalvincent_tensor(path, strict=False)
+    if x.dtype != torch.uint8:
+        raise TypeError(f"x should be of dtype torch.uint8 instead of {x.dtype}")
+    if x.ndimension() != 1:
+        raise ValueError(f"x should have 1 dimension instead of {x.ndimension()}")
+    return x.long()
+
+
+def read_image_file(path: str) -> torch.Tensor:
+    x = read_sn3_pascalvincent_tensor(path, strict=False)
+    if x.dtype != torch.uint8:
+        raise TypeError(f"x should be of dtype torch.uint8 instead of {x.dtype}")
+    if x.ndimension() != 3:
+        raise ValueError(f"x should have 3 dimension instead of {x.ndimension()}")
+    return x

MTIL_datasets/oxford_flowers.py

@@ -0,0 +1,89 @@
+import os
+import pickle
+import random
+from scipy.io import loadmat
+from collections import defaultdict
+
+from .utils import *
+
+from .oxford_pets import OxfordPets
+
+
+class OxfordFlowers(DatasetBase):
+
+    dataset_dir = "oxford_flowers"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.image_dir = os.path.join(self.dataset_dir, "jpg")
+        self.label_file = os.path.join(self.dataset_dir, "imagelabels.mat")
+        self.lab2cname_file = os.path.join(self.dataset_dir, "cat_to_name.json")
+        self.split_path = os.path.join(self.dataset_dir, "split_zhou_OxfordFlowers.json")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        if os.path.exists(self.split_path):
+            train, val, test = OxfordPets.read_split(self.split_path, self.image_dir)
+        else:
+            train, val, test = self.read_data()
+            OxfordPets.save_split(train, val, test, self.split_path, self.image_dir)
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        subsample = subsample_classes
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample)
+
+        self.templates = [
+            lambda c: f"a photo of a {c}, a type of flower.",
+        ]
+
+        super().__init__(train_x=train, val=val, test=test)
+
+    def read_data(self):
+        tracker = defaultdict(list)
+        label_file = loadmat(self.label_file)["labels"][0]
+        for i, label in enumerate(label_file):
+            imname = f"image_{str(i + 1).zfill(5)}.jpg"
+            impath = os.path.join(self.image_dir, imname)
+            label = int(label)
+            tracker[label].append(impath)
+
+        print("Splitting data into 50% train, 20% val, and 30% test")
+
+        def _collate(ims, y, c):
+            items = []
+            for im in ims:
+                item = Datum(impath=im, label=y - 1, classname=c)
+                items.append(item)
+            return items
+
+        lab2cname = read_json(self.lab2cname_file)
+        train, val, test = [], [], []
+        for label, impaths in tracker.items():
+            random.shuffle(impaths)
+            n_total = len(impaths)
+            n_train = round(n_total * 0.5)
+            n_val = round(n_total * 0.2)
+            n_test = n_total - n_train - n_val
+            assert n_train > 0 and n_val > 0 and n_test > 0
+            cname = lab2cname[str(label)]
+            train.extend(_collate(impaths[:n_train], label, cname))
+            val.extend(_collate(impaths[n_train : n_train + n_val], label, cname))
+            test.extend(_collate(impaths[n_train + n_val :], label, cname))
+
+        return train, val, test

MTIL_datasets/oxford_pets.py

@@ -0,0 +1,176 @@
+import os
+import pickle
+import math
+import random
+from collections import defaultdict
+
+from .utils import *
+
+
+class OxfordPets(DatasetBase):
+
+    dataset_dir = "oxford_pets"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.image_dir = os.path.join(self.dataset_dir, "images")
+        self.anno_dir = os.path.join(self.dataset_dir, "annotations")
+        self.split_path = os.path.join(self.dataset_dir, "split_zhou_OxfordPets.json")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        if os.path.exists(self.split_path):
+            train, val, test = self.read_split(self.split_path, self.image_dir)
+        else:
+            trainval = self.read_data(split_file="trainval.txt")
+            test = self.read_data(split_file="test.txt")
+            train, val = self.split_trainval(trainval)
+            self.save_split(train, val, test, self.split_path, self.image_dir)
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        subsample = subsample_classes
+        train, val, test = self.subsample_classes(train, val, test, subsample=subsample)
+
+        self.templates = [
+            lambda c: f"a photo of a {c}, a type of pet.",
+        ]
+
+        super().__init__(train_x=train, val=val, test=test)
+
+    def read_data(self, split_file):
+        filepath = os.path.join(self.anno_dir, split_file)
+        items = []
+
+        with open(filepath, "r") as f:
+            lines = f.readlines()
+            for line in lines:
+                line = line.strip()
+                imname, label, species, _ = line.split(" ")
+                breed = imname.split("_")[:-1]
+                breed = "_".join(breed)
+                breed = breed.lower()
+                imname += ".jpg"
+                impath = os.path.join(self.image_dir, imname)
+                label = int(label) - 1  # convert to 0-based index
+                item = Datum(impath=impath, label=label, classname=breed)
+                items.append(item)
+
+        return items
+
+    @staticmethod
+    def split_trainval(trainval, p_val=0.2):
+        p_trn = 1 - p_val
+        tracker = defaultdict(list)
+        for idx, item in enumerate(trainval):
+            label = item.label
+            tracker[label].append(idx)
+
+        train, val = [], []
+        for label, idxs in tracker.items():
+            n_val = round(len(idxs) * p_val)
+            assert n_val > 0
+            random.shuffle(idxs)
+            for n, idx in enumerate(idxs):
+                item = trainval[idx]
+                if n < n_val:
+                    val.append(item)
+                else:
+                    train.append(item)
+
+        return train, val
+
+    @staticmethod
+    def save_split(train, val, test, filepath, path_prefix):
+        def _extract(items):
+            out = []
+            for item in items:
+                impath = item.impath
+                label = item.label
+                classname = item.classname
+                impath = impath.replace(path_prefix, "")
+                if impath.startswith("/"):
+                    impath = impath[1:]
+                out.append((impath, label, classname))
+            return out
+
+        train = _extract(train)
+        val = _extract(val)
+        test = _extract(test)
+
+        split = {"train": train, "val": val, "test": test}
+
+        write_json(split, filepath)
+        print(f"Saved split to {filepath}")
+
+    @staticmethod
+    def read_split(filepath, path_prefix):
+        def _convert(items):
+            out = []
+            for impath, label, classname in items:
+                impath = os.path.join(path_prefix, impath)
+                item = Datum(impath=impath, label=int(label), classname=classname)
+                out.append(item)
+            return out
+
+        print(f"Reading split from {filepath}")
+        split = read_json(filepath)
+        train = _convert(split["train"])
+        val = _convert(split["val"])
+        test = _convert(split["test"])
+
+        return train, val, test
+    
+    @staticmethod
+    def subsample_classes(*args, subsample="all"):
+        assert subsample in ["all", "base", "new"]
+
+        if subsample == "all":
+            return args
+        
+        dataset = args[0]
+        labels = set()
+        for item in dataset:
+            labels.add(item.label)
+        labels = list(labels)
+        labels.sort()
+        n = len(labels)
+        m = math.ceil(n / 2)
+
+        print(f"SUBSAMPLE {subsample.upper()} CLASSES!")
+        if subsample == "base":
+            selected = labels[:m]
+        else:
+            selected = labels[m:]
+        relabeler = {y: y_new for y_new, y in enumerate(selected)}
+        
+        output = []
+        for dataset in args:
+            dataset_new = []
+            for item in dataset:
+                if item.label not in selected:
+                    continue
+                item_new = Datum(
+                    impath=item.impath,
+                    label=relabeler[item.label],
+                    classname=item.classname
+                )
+                dataset_new.append(item_new)
+            output.append(dataset_new)
+        
+        return output

MTIL_datasets/pcam.py

@@ -0,0 +1,198 @@
+import os
+import pickle
+import random
+from typing import List, Tuple
+
+import h5py
+import numpy as np
+
+from .utils import Datum, DatasetBase, mkdir_if_missing, read_json, write_json
+from .oxford_pets import OxfordPets
+
+
+PCAM_CLASSES = [
+    'lymph node tissue without metastatic tumor',
+    'metastatic tumor in lymph node tissue',
+]
+
+# Multiple domain-specific templates for histopathology microscopy images
+PCAM_TEMPLATES = [
+    'a microscopy image patch of {}',
+    'a histopathology image of {}',
+    'a hematoxylin and eosin stained image of {}',
+    'a high-resolution histology patch of {}',
+    'a digital pathology slide patch of {}',
+    'this is a microscopy image of {}',
+    'this is a histopathology image of {}',
+]
+
+
+class PCam(DatasetBase):
+
+    dataset_dir = 'pcam'
+
+    def __init__(self, root, num_shots: int = 0, seed: int = 1, subsample_classes: str = 'all', val_ratio: float = 0.2):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        mkdir_if_missing(self.dataset_dir)
+
+        # HDF5 file paths
+        self.train_x_path = os.path.join(self.dataset_dir, 'camelyonpatch_level_2_split_train_x.h5')
+        self.train_y_path = os.path.join(self.dataset_dir, 'camelyonpatch_level_2_split_train_y.h5')
+        self.test_x_path = os.path.join(self.dataset_dir, 'camelyonpatch_level_2_split_test_x.h5')
+        self.test_y_path = os.path.join(self.dataset_dir, 'camelyonpatch_level_2_split_test_y.h5')
+
+        for p in [self.train_x_path, self.train_y_path, self.test_x_path, self.test_y_path]:
+            if not os.path.isfile(p):
+                raise FileNotFoundError(f"PCam: expected file not found: {p}")
+
+        self.split_path = os.path.join(self.dataset_dir, 'split_custom_PCam.json')
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, 'split_fewshot')
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        random.seed(seed)
+        np.random.seed(seed)
+
+        if os.path.exists(self.split_path):
+            train, val, test = self.read_split(self.split_path)
+        else:
+            trainval = self._read_train(self.train_x_path, self.train_y_path)
+            test = self._read_test(self.test_x_path, self.test_y_path)
+            train, val = self.split_trainval(trainval, p_val=val_ratio)
+            self.save_split(train, val, test, self.split_path)
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, 'rb') as f:
+                    data = pickle.load(f)
+                    train, val = data['train'], data['val']
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {'train': train, 'val': val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, 'wb') as f:
+                    pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
+
+        # Optional class subsampling (kept for consistency with other datasets)
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample_classes)
+
+        # Templates per user specification
+        self.templates = PCAM_TEMPLATES
+
+        # Debug: print split sizes and label histograms
+        def _hist(items):
+            from collections import Counter
+            cnt = Counter([it.label for it in items])
+            # ensure keys present for both binary classes
+            out = {i: int(cnt.get(i, 0)) for i in range(len(PCAM_CLASSES))}
+            return out
+
+        super().__init__(train_x=train, val=val, test=test)
+        # Ensure stable binary classification metadata
+        self._classnames = PCAM_CLASSES
+        self._lab2cname = {i: c for i, c in enumerate(PCAM_CLASSES)}
+        self._num_classes = len(PCAM_CLASSES)
+
+    @staticmethod
+    def _first_key(h5_path: str) -> str:
+        with h5py.File(h5_path, 'r') as f:
+            keys = list(f.keys())
+            if not keys:
+                raise RuntimeError(f"No datasets found in H5 file: {h5_path}")
+            return keys[0]
+
+    @staticmethod
+    def _read_labels(h5_path: str) -> np.ndarray:
+        key = PCam._first_key(h5_path)
+        with h5py.File(h5_path, 'r') as f:
+            y = f[key][...]
+        y = np.asarray(y).squeeze()
+        y = y.astype(np.int64)
+        return y
+
+    def _read_train(self, x_path: str, y_path: str) -> List[Datum]:
+        x_key = self._first_key(x_path)
+        y = self._read_labels(y_path)
+        items: List[Datum] = []
+        for i, label in enumerate(y.tolist()):
+            label_i = int(label)
+            classname = PCAM_CLASSES[label_i]
+            # Lazy image reference: ('h5', abs_h5_path, dataset_key, index)
+            abs_path = os.path.abspath(x_path)
+            impath = ('h5', abs_path, x_key, i)
+            items.append(Datum(impath=impath, label=label_i, classname=classname))
+        return items
+
+    def _read_test(self, x_path: str, y_path: str) -> List[Datum]:
+        x_key = self._first_key(x_path)
+        y = self._read_labels(y_path)
+        items: List[Datum] = []
+        for i, label in enumerate(y.tolist()):
+            label_i = int(label)
+            classname = PCAM_CLASSES[label_i]
+            abs_path = os.path.abspath(x_path)
+            impath = ('h5', abs_path, x_key, i)
+            items.append(Datum(impath=impath, label=label_i, classname=classname))
+        return items
+
+    @staticmethod
+    def split_trainval(trainval: List[Datum], p_val: float = 0.2) -> Tuple[List[Datum], List[Datum]]:
+        from collections import defaultdict
+        p_trn = 1 - p_val
+        print(f"Splitting PCam train into {p_trn:.0%} train and {p_val:.0%} val")
+        tracker = defaultdict(list)
+        for idx, item in enumerate(trainval):
+            tracker[item.label].append(idx)
+        train, val = [], []
+        for _, idxs in tracker.items():
+            n_val = max(1, round(len(idxs) * p_val))
+            random.shuffle(idxs)
+            for n, i in enumerate(idxs):
+                if n < n_val:
+                    val.append(trainval[i])
+                else:
+                    train.append(trainval[i])
+        return train, val
+
+    def save_split(self, train: List[Datum], val: List[Datum], test: List[Datum], filepath: str):
+        def _ser(items: List[Datum]):
+            out = []
+            for it in items:
+                impath = it.impath
+                if isinstance(impath, tuple) and len(impath) == 4 and impath[0] == 'h5':
+                    tag, abs_path, key, idx = impath
+                    # store relative path for portability
+                    rel = os.path.relpath(abs_path, self.dataset_dir)
+                    impath_ser = [tag, rel, key, int(idx)]
+                else:
+                    raise ValueError('PCam expects H5 tuple paths')
+                out.append((impath_ser, int(it.label), it.classname))
+            return out
+        split = {
+            'train': _ser(train),
+            'val': _ser(val),
+            'test': _ser(test),
+        }
+        write_json(split, filepath)
+        print(f"Saved PCam split to {filepath}")
+
+    def read_split(self, filepath: str):
+        def _deser(items):
+            out = []
+            for impath_ser, label, classname in items:
+                if isinstance(impath_ser, (list, tuple)) and len(impath_ser) == 4 and impath_ser[0] == 'h5':
+                    tag, rel, key, idx = impath_ser
+                    fpath = os.path.join(self.dataset_dir, rel)
+                    impath = (tag, fpath, key, int(idx))
+                else:
+                    raise ValueError('PCam split contains invalid path entries')
+                out.append(Datum(impath=impath, label=int(label), classname=classname))
+            return out
+        split = read_json(filepath)
+        train = _deser(split['train'])
+        val = _deser(split['val'])
+        test = _deser(split['test'])
+        return train, val, test

MTIL_datasets/resisc.py

@@ -0,0 +1,199 @@
+import os
+import pickle
+import random
+import warnings
+
+from .utils import *
+from .oxford_pets import OxfordPets
+
+# Canonical class list for RESISC45 (order defines label ids 0..44)
+RESISC_CLASSES = [
+    'airplane',
+    'airport',
+    'baseball diamond',
+    'basketball court',
+    'beach',
+    'bridge',
+    'chaparral',
+    'church',
+    'circular farmland',
+    'cloud',
+    'commercial area',
+    'dense residential',
+    'desert',
+    'forest',
+    'freeway',
+    'golf course',
+    'ground track field',
+    'harbor',
+    'industrial area',
+    'intersection',
+    'island',
+    'lake',
+    'meadow',
+    'medium residential',
+    'mobile home park',
+    'mountain',
+    'overpass',
+    'palace',
+    'parking lot',
+    'railway',
+    'railway station',
+    'rectangular farmland',
+    'river',
+    'roundabout',
+    'runway',
+    'sea ice',
+    'ship',
+    'snowberg',
+    'sparse residential',
+    'stadium',
+    'storage tank',
+    'tennis court',
+    'terrace',
+    'thermal power station',
+    'wetland',
+]
+
+# Prompt templates (strings) as provided
+RESISC_TEMPLATES = [
+    'satellite imagery of {}.',
+    'aerial imagery of {}.',
+    'satellite photo of {}.',
+    'aerial photo of {}.',
+    'satellite view of {}.',
+    'aerial view of {}.',
+    'satellite imagery of a {}.',
+    'aerial imagery of a {}.',
+    'satellite photo of a {}.',
+    'aerial photo of a {}.',
+    'satellite view of a {}.',
+    'aerial view of a {}.',
+    'satellite imagery of the {}.',
+    'aerial imagery of the {}.',
+    'satellite photo of the {}.',
+    'aerial photo of the {}.',
+    'satellite view of the {}.',
+    'aerial view of the {}.',
+]
+
+RESISC_DEBUG = os.environ.get("RESISC_DEBUG", "0") not in ("0", "false", "False", "")
+
+def _dbg(msg: str):
+    if RESISC_DEBUG:
+        print(f"[RESISC45][DEBUG] {msg}")
+
+
+def _norm_name(s: str) -> str:
+    s = s.lower().strip()
+    for ch in [" ", "_", "-", "."]:
+        s = s.replace(ch, "")
+    return s
+
+
+class RESISC45(DatasetBase):
+
+    dataset_dir = "resisc45"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.image_dir = self.dataset_dir  # images are under class subfolders directly
+        self.split_path = os.path.join(self.dataset_dir, "split_custom_RESISC45.json")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+        _dbg(f"dataset_dir={self.dataset_dir}")
+
+        if os.path.exists(self.split_path):
+            try:
+                train, val, test = OxfordPets.read_split(self.split_path, self.image_dir)
+            except Exception as e:
+                warnings.warn(f"RESISC45: failed to read split file '{self.split_path}'; rebuilding. Error: {e}")
+                train, val, test = self.read_and_split_data(self.image_dir)
+                try:
+                    OxfordPets.save_split(train, val, test, self.split_path, self.image_dir)
+                except Exception as e2:
+                    warnings.warn(f"RESISC45: failed to save rebuilt split to '{self.split_path}': {e2}")
+        else:
+            train, val, test = self.read_and_split_data(self.image_dir)
+            try:
+                OxfordPets.save_split(train, val, test, self.split_path, self.image_dir)
+            except Exception as e:
+                warnings.warn(f"RESISC45: failed to save split to '{self.split_path}': {e}")
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample_classes)
+        self.templates = RESISC_TEMPLATES
+        super().__init__(train_x=train, val=val, test=test)
+
+    @staticmethod
+    def read_and_split_data(image_dir, p_trn=0.5, p_val=0.2):
+        try:
+            categories = listdir_nohidden(image_dir, sort=True)
+        except Exception as e:
+            warnings.warn(f"RESISC45: failed to list directory '{image_dir}': {e}")
+            raise
+        categories = [c for c in categories if os.path.isdir(os.path.join(image_dir, c))]
+        cat_norm = {_norm_name(c): c for c in categories}
+
+        missing = []
+        class_to_dir = {}
+        for cname in RESISC_CLASSES:
+            key = _norm_name(cname)
+            if key in cat_norm:
+                class_to_dir[cname] = cat_norm[key]
+            else:
+                missing.append(cname)
+        if missing:
+            warnings.warn(f"RESISC45: missing class folders for {len(missing)} classes: {missing[:5]}{' ...' if len(missing)>5 else ''}")
+
+        def _collate(paths, y, cname):
+            return [Datum(impath=p, label=y, classname=cname) for p in paths]
+
+        train, val, test = [], [], []
+        for y, cname in enumerate(RESISC_CLASSES):
+            if cname not in class_to_dir:
+                continue
+            cdir = os.path.join(image_dir, class_to_dir[cname])
+            try:
+                images = listdir_nohidden(cdir, sort=False)
+            except Exception as e:
+                warnings.warn(f"RESISC45: failed to list class folder '{cdir}': {e}")
+                images = []
+            images = [os.path.join(cdir, im) for im in images]
+            random.shuffle(images)
+            n_total = len(images)
+            if n_total == 0:
+                warnings.warn(f"RESISC45: empty class folder {cdir}; skipping")
+                continue
+            n_train = round(n_total * p_trn)
+            n_val = round(n_total * p_val)
+            n_test = n_total - n_train - n_val
+            if not (n_train > 0 and n_val > 0 and n_test > 0):
+                # Fallback to keep all splits non-empty where possible
+                if n_total >= 3:
+                    n_train, n_val, n_test = 1, 1, n_total - 2
+                elif n_total == 2:
+                    n_train, n_val, n_test = 1, 1, 0
+                else:  # 1
+                    n_train, n_val, n_test = 1, 0, 0
+
+            train.extend(_collate(images[:n_train], y, cname))
+            val.extend(_collate(images[n_train:n_train + n_val], y, cname))
+            test.extend(_collate(images[n_train + n_val:], y, cname))
+
+        return train, val, test

MTIL_datasets/sst2.py

@@ -0,0 +1,122 @@
+import os
+import pickle
+
+from .utils import *
+from .oxford_pets import OxfordPets
+from .dtd import DescribableTextures as DTD
+
+
+class SST2(DatasetBase):
+    """
+    SST2 (rendered) dataset loader for MTIL.
+
+    Expected structure (either of the following roots):
+      <root>/sst2/
+        ├─ train/{negative,positive}/*.png
+        ├─ valid/{negative,positive}/*.png
+        └─ test/{negative,positive}/*.png
+
+      <root>/rendered-sst2/  (alias supported)
+        ├─ train/{negative,positive}/*.png
+        ├─ valid/{negative,positive}/*.png
+        └─ test/{negative,positive}/*.png
+
+    Classes: ['negative', 'positive']
+    Templates: ['a {} review of a movie.']
+    """
+
+    dataset_dir = "sst2"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        primary_dir = os.path.join(root, self.dataset_dir)
+        alt_dir = os.path.join(root, "rendered-sst2")
+
+        if os.path.isdir(primary_dir):
+            self.dataset_dir = primary_dir
+        elif os.path.isdir(alt_dir):
+            self.dataset_dir = alt_dir
+        else:
+            raise ValueError(
+                "SST2: dataset folder not found. Expected one of: '{}' or '{}'".format(primary_dir, alt_dir)
+            )
+
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        train_dir = os.path.join(self.dataset_dir, "train")
+        valid_dir = os.path.join(self.dataset_dir, "valid")
+        test_dir = os.path.join(self.dataset_dir, "test")
+
+        use_folder_splits = os.path.isdir(train_dir) and os.path.isdir(valid_dir) and os.path.isdir(test_dir)
+
+        # fixed class order and validation
+        classes = ["negative", "positive"]
+        class_to_label = {c: i for i, c in enumerate(classes)}
+
+        if use_folder_splits:
+            train = self._read_split_dir(train_dir, class_to_label)
+            val = self._read_split_dir(valid_dir, class_to_label)
+            test = self._read_split_dir(test_dir, class_to_label)
+        else:
+            # Fallbacks consistent with other datasets: look for a JSON split, else naive split
+            image_dir = os.path.join(self.dataset_dir, "images")
+            self.image_dir = image_dir if os.path.isdir(image_dir) else self.dataset_dir
+            split_path = os.path.join(self.dataset_dir, "split_zhou_SST2.json")
+            if os.path.exists(split_path):
+                train, val, test = OxfordPets.read_split(split_path, self.image_dir)
+            else:
+                train, val, test = DTD.read_and_split_data(self.image_dir)
+                OxfordPets.save_split(train, val, test, split_path, self.image_dir)
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        subsample = subsample_classes
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample)
+
+        # Template per user spec
+        self.templates = [
+            lambda c: f'a {c} review of a movie.',
+        ]
+
+        super().__init__(train_x=train, val=val, test=test)
+
+    def _read_split_dir(self, split_dir, class_to_label):
+        items = []
+        if not os.path.isdir(split_dir):
+            return items
+        codes = listdir_nohidden(split_dir)
+        # Validate folders: must be subset of expected classes and cover at least one class
+        unexpected = sorted([c for c in codes if c not in class_to_label])
+        if unexpected:
+            raise ValueError(
+                f"SST2: Found unexpected class folders in '{split_dir}': {unexpected}. "
+                f"Expected only {list(class_to_label.keys())}."
+            )
+        for cls in codes:
+            class_dir = os.path.join(split_dir, cls)
+            if not os.path.isdir(class_dir):
+                continue
+            label = class_to_label.get(cls)
+            if label is None:
+                raise ValueError(
+                    f"SST2: Inconsistent label mapping for class '{cls}' in split '{split_dir}'."
+                )
+            cname = cls
+            for fname in listdir_nohidden(class_dir):
+                impath = os.path.join(class_dir, fname)
+                items.append(Datum(impath=impath, label=label, classname=cname))
+        return items

MTIL_datasets/stanford_cars.py

@@ -0,0 +1,83 @@
+import os
+import pickle
+from scipy.io import loadmat
+
+from .utils import *
+
+from .oxford_pets import OxfordPets
+
+
+class StanfordCars(DatasetBase):
+
+    dataset_dir = "stanford_cars"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.split_path = os.path.join(self.dataset_dir, "split_zhou_StanfordCars.json")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        if os.path.exists(self.split_path):
+            train, val, test = OxfordPets.read_split(self.split_path, self.dataset_dir)
+        else:
+            trainval_file = os.path.join(self.dataset_dir, "devkit", "cars_train_annos.mat")
+            test_file = os.path.join(self.dataset_dir, "cars_test_annos_withlabels.mat")
+            meta_file = os.path.join(self.dataset_dir, "devkit", "cars_meta.mat")
+            trainval = self.read_data("cars_train", trainval_file, meta_file)
+            test = self.read_data("cars_test", test_file, meta_file)
+            train, val = OxfordPets.split_trainval(trainval)
+            OxfordPets.save_split(train, val, test, self.split_path, self.dataset_dir)
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        subsample = subsample_classes
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample)
+
+        self.templates = [
+            lambda c: f"a photo of a {c}, a type of car.",
+            lambda c: f"a photo of a {c}.",
+            lambda c: f"a photo of the {c}.",
+            lambda c: f"a photo of my {c}.",
+            lambda c: f"i love my {c}!",
+            lambda c: f"a photo of my dirty {c}.",
+            lambda c: f"a photo of my clean {c}.",
+            lambda c: f"a photo of my new {c}.",
+            lambda c: f"a photo of my old {c}.",
+        ]
+
+        super().__init__(train_x=train, val=val, test=test)
+
+    def read_data(self, image_dir, anno_file, meta_file):
+        anno_file = loadmat(anno_file)["annotations"][0]
+        meta_file = loadmat(meta_file)["class_names"][0]
+        items = []
+
+        for i in range(len(anno_file)):
+            imname = anno_file[i]["fname"][0]
+            impath = os.path.join(self.dataset_dir, image_dir, imname)
+            label = anno_file[i]["class"][0, 0]
+            label = int(label) - 1
+            classname = meta_file[label][0]
+            names = classname.split(" ")
+            year = names.pop(-1)
+            names.insert(0, year)
+            classname = " ".join(names)
+            item = Datum(impath=impath, label=label, classname=classname)
+            items.append(item)
+
+        return items

MTIL_datasets/stl10.py

@@ -0,0 +1,143 @@
+import os
+import math
+import random
+from typing import List, Tuple
+
+from .utils import *  # Datum, DatasetBase, listdir_nohidden
+from .oxford_pets import OxfordPets
+
+STL10_CLASSES: List[str] = [
+    'airplane',
+    'bird',
+    'car',
+    'cat',
+    'deer',
+    'dog',
+    'horse',
+    'monkey',
+    'ship',
+    'truck',
+]
+
+# keep templates as strings with {}
+STL10_TEMPLATES: List[str] = [
+    'a photo of a {}.',
+    'a photo of the {}.',
+]
+
+
+class STL10(DatasetBase):
+
+    dataset_dir = "stl10"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all', test_ratio: float = 0.2):
+        """
+        Expect directory structure:
+            {root}/stl10/<class_name>/*.png|jpg|jpeg|bmp|webp
+        No official test set provided; we split per-class into train/test.
+        Then we further split train into train/val using OxfordPets.split_trainval.
+        """
+        rnd = random.Random(seed)
+
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.split_path = os.path.join(self.dataset_dir, "split_custom_STL10.json")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        # If we have a saved split, reuse for reproducibility
+        if os.path.exists(self.split_path):
+            train, val, test = OxfordPets.read_split(self.split_path, self.dataset_dir)
+        else:
+            trainval, test = self._read_from_folders(self.dataset_dir, rnd=rnd, test_ratio=test_ratio)
+            train, val = self._split_trainval_safe(trainval)
+            OxfordPets.save_split(train, val, test, self.split_path, self.dataset_dir)
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                import pickle
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                import pickle
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample_classes)
+
+        self.templates = STL10_TEMPLATES
+
+        super().__init__(train_x=train, val=val, test=test)
+
+    def _read_from_folders(self, data_dir: str, rnd: random.Random, test_ratio: float) -> Tuple[List[Datum], List[Datum]]:
+        exts = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}
+        class_to_label = {name: i for i, name in enumerate(STL10_CLASSES)}
+
+        items_by_label = {i: [] for i in range(len(STL10_CLASSES))}
+        for cname in STL10_CLASSES:
+            cdir = os.path.join(data_dir, cname)
+            if not os.path.isdir(cdir):
+                raise FileNotFoundError(f"Class folder not found: {cdir}")
+            for fname in listdir_nohidden(cdir, sort=True):
+                fext = os.path.splitext(fname)[1].lower()
+                if fext not in exts:
+                    continue
+                impath = os.path.join(cdir, fname)
+                label = class_to_label[cname]
+                items_by_label[label].append(Datum(impath=impath, label=label, classname=cname))
+
+        trainval, test = [], []
+        for label, items in items_by_label.items():
+            if not items:
+                continue
+            rnd.shuffle(items)
+            if len(items) == 1:
+                # keep the only sample for training to preserve class presence in train
+                trainval.extend(items)
+                continue
+            n_test = max(1, int(round(len(items) * test_ratio)))
+            if n_test >= len(items):
+                n_test = len(items) - 1
+            test.extend(items[:n_test])
+            trainval.extend(items[n_test:])
+
+        return trainval, test
+
+    @staticmethod
+    def _split_trainval_safe(trainval: List[Datum], p_val: float = 0.2) -> Tuple[List[Datum], List[Datum]]:
+        from collections import defaultdict
+        tracker = defaultdict(list)
+        for idx, item in enumerate(trainval):
+            tracker[item.label].append(idx)
+
+        train, val = [], []
+        for label, idxs in tracker.items():
+            n = len(idxs)
+            if n <= 1:
+                # not enough to create a val sample for this class
+                for idx in idxs:
+                    train.append(trainval[idx])
+                continue
+            n_val = max(1, int(round(n * p_val)))
+            if n_val >= n:
+                n_val = n - 1
+            random.shuffle(idxs)
+            for i, idx in enumerate(idxs):
+                if i < n_val:
+                    val.append(trainval[idx])
+                else:
+                    train.append(trainval[idx])
+
+        # If val ended up empty (degenerate tiny dataset), move one from train to val
+        if len(val) == 0 and len(train) > 0:
+            val.append(train[-1])
+            train = train[:-1]
+
+        return train, val

MTIL_datasets/sun397.py

@@ -0,0 +1,81 @@
+import os
+import pickle
+
+from .utils import *
+
+from .oxford_pets import OxfordPets
+
+
+class SUN397(DatasetBase):
+
+    dataset_dir = "sun397"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.image_dir = os.path.join(self.dataset_dir, "SUN397")
+        self.split_path = os.path.join(self.dataset_dir, "split_zhou_SUN397.json")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        if os.path.exists(self.split_path):
+            train, val, test = OxfordPets.read_split(self.split_path, self.image_dir)
+        else:
+            classnames = []
+            with open(os.path.join(self.dataset_dir, "ClassName.txt"), "r") as f:
+                lines = f.readlines()
+                for line in lines:
+                    line = line.strip()[1:]
+                    classnames.append(line)
+            cname2lab = {c: i for i, c in enumerate(classnames)}
+            trainval = self.read_data(cname2lab, "Training_01.txt")
+            test = self.read_data(cname2lab, "Testing_01.txt")
+            train, val = OxfordPets.split_trainval(trainval)
+            OxfordPets.save_split(train, val, test, self.split_path, self.image_dir)
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        subsample = subsample_classes
+        train, val, test = OxfordPets.subsample_classes(train, val, test, subsample=subsample)
+
+        self.templates = [
+            lambda c: f"a photo of a {c}.",
+            lambda c: f"a photo of the {c}.",
+        ]
+
+        super().__init__(train_x=train, val=val, test=test)
+
+    def read_data(self, cname2lab, text_file):
+        text_file = os.path.join(self.dataset_dir, text_file)
+        items = []
+
+        with open(text_file, "r") as f:
+            lines = f.readlines()
+            for line in lines:
+                imname = line.strip()[1:]
+                classname = os.path.dirname(imname)
+                label = cname2lab[classname]
+                impath = os.path.join(self.image_dir, imname)
+
+                names = classname.split("/")[1:]
+                names = names[::-1]
+                classname = " ".join(names)
+
+                item = Datum(impath=impath, label=label, classname=classname)
+                items.append(item)
+
+        return items

MTIL_datasets/ucf101.py

@@ -0,0 +1,360 @@
+import os
+import pickle
+import random
+import math
+from collections import defaultdict
+from typing import List, Tuple
+
+from .utils import *  # Datum, DatasetBase, mkdir_if_missing, read_json, write_json, listdir_nohidden
+
+
+# Canonical class names as provided
+UCF101_CANONICAL: List[str] = [
+    'Apply Eye Makeup',
+    'Apply Lipstick',
+    'Archery',
+    'Baby Crawling',
+    'Balance Beam',
+    'Band Marching',
+    'Baseball Pitch',
+    'Basketball',
+    'Basketball Dunk',
+    'Bench Press',
+    'Biking',
+    'Billiards',
+    'Blow Dry Hair',
+    'Blowing Candles',
+    'Body Weight Squats',
+    'Bowling',
+    'Boxing Punching Bag',
+    'Boxing Speed Bag',
+    'Breast Stroke',
+    'Brushing Teeth',
+    'Clean And Jerk',
+    'Cliff Diving',
+    'Cricket Bowling',
+    'Cricket Shot',
+    'Cutting In Kitchen',
+    'Diving',
+    'Drumming',
+    'Fencing',
+    'Field Hockey Penalty',
+    'Floor Gymnastics',
+    'Frisbee Catch',
+    'Front Crawl',
+    'Golf Swing',
+    'Haircut',
+    'Hammer Throw',
+    'Hammering',
+    'Hand Stand Pushups',
+    'Handstand Walking',
+    'Head Massage',
+    'High Jump',
+    'Horse Race',
+    'Horse Riding',
+    'Hula Hoop',
+    'Ice Dancing',
+    'Javelin Throw',
+    'Juggling Balls',
+    'Jump Rope',
+    'Jumping Jack',
+    'Kayaking',
+    'Knitting',
+    'Long Jump',
+    'Lunges',
+    'Military Parade',
+    'Mixing',
+    'Mopping Floor',
+    'Nunchucks',
+    'Parallel Bars',
+    'Pizza Tossing',
+    'Playing Cello',
+    'Playing Daf',
+    'Playing Dhol',
+    'Playing Flute',
+    'Playing Guitar',
+    'Playing Piano',
+    'Playing Sitar',
+    'Playing Tabla',
+    'Playing Violin',
+    'Pole Vault',
+    'Pommel Horse',
+    'Pull Ups',
+    'Punch',
+    'Push Ups',
+    'Rafting',
+    'Rock Climbing Indoor',
+    'Rope Climbing',
+    'Rowing',
+    'Salsa Spin',
+    'Shaving Beard',
+    'Shotput',
+    'Skate Boarding',
+    'Skiing',
+    'Skijet',
+    'Sky Diving',
+    'Soccer Juggling',
+    'Soccer Penalty',
+    'Still Rings',
+    'Sumo Wrestling',
+    'Surfing',
+    'Swing',
+    'Table Tennis Shot',
+    'Tai Chi',
+    'Tennis Swing',
+    'Throw Discus',
+    'Trampoline Jumping',
+    'Typing',
+    'Uneven Bars',
+    'Volleyball Spiking',
+    'Walking With Dog',
+    'Wall Pushups',
+    'Writing On Board',
+    'Yo Yo',
+]
+
+
+UCF101_TEMPLATES: List[str] = [
+    'a photo of a person {}.',
+    'a video of a person {}.',
+    'a example of a person {}.',
+    'a demonstration of a person {}.',
+    'a photo of the person {}.',
+    'a video of the person {}.',
+    'a example of the person {}.',
+    'a demonstration of the person {}.',
+    'a photo of a person using {}.',
+    'a video of a person using {}.',
+    'a example of a person using {}.',
+    'a demonstration of a person using {}.',
+    'a photo of the person using {}.',
+    'a video of the person using {}.',
+    'a example of the person using {}.',
+    'a demonstration of the person using {}.',
+    'a photo of a person doing {}.',
+    'a video of a person doing {}.',
+    'a example of a person doing {}.',
+    'a demonstration of a person doing {}.',
+    'a photo of the person doing {}.',
+    'a video of the person doing {}.',
+    'a example of the person doing {}.',
+    'a demonstration of the person doing {}.',
+    'a photo of a person during {}.',
+    'a video of a person during {}.',
+    'a example of a person during {}.',
+    'a demonstration of a person during {}.',
+    'a photo of the person during {}.',
+    'a video of the person during {}.',
+    'a example of the person during {}.',
+    'a demonstration of the person during {}.',
+    'a photo of a person performing {}.',
+    'a video of a person performing {}.',
+    'a example of a person performing {}.',
+    'a demonstration of a person performing {}.',
+    'a photo of the person performing {}.',
+    'a video of the person performing {}.',
+    'a example of the person performing {}.',
+    'a demonstration of the person performing {}.',
+    'a photo of a person practicing {}.',
+    'a video of a person practicing {}.',
+    'a example of a person practicing {}.',
+    'a demonstration of a person practicing {}.',
+    'a photo of the person practicing {}.',
+    'a video of the person practicing {}.',
+    'a example of the person practicing {}.',
+    'a demonstration of the person practicing {}.',
+]
+
+
+class UCF101(DatasetBase):
+    """
+    UCF101 midframes classification dataset adapter for MTIL.
+    Expected structure:
+      root/ucf101/
+        UCF-101-midframes/<Class_Folder>/<image>.jpg
+        split_zhou_UCF101.json
+    """
+
+    dataset_dir = "ucf101"
+
+    def __init__(self, root, num_shots=0, seed=1, subsample_classes='all'):
+        root = os.path.abspath(os.path.expanduser(root))
+        self.dataset_dir = os.path.join(root, self.dataset_dir)
+        self.image_dir = os.path.join(self.dataset_dir, "UCF-101-midframes")
+        self.split_path = os.path.join(self.dataset_dir, "split_zhou_UCF101.json")
+        self.split_fewshot_dir = os.path.join(self.dataset_dir, "split_fewshot")
+        mkdir_if_missing(self.split_fewshot_dir)
+
+        if os.path.exists(self.split_path):
+            train, val, test = self.read_split(self.split_path, self.image_dir)
+        else:
+            # Fallback: build from directory and split train/val; use val as test too.
+            trainval = self._read_from_dir()
+            train, val = self.split_trainval(trainval)
+            test = list(val)
+            self.save_split(train, val, test, self.split_path, self.image_dir)
+
+        if num_shots >= 1:
+            preprocessed = os.path.join(self.split_fewshot_dir, f"shot_{num_shots}-seed_{seed}.pkl")
+            if os.path.exists(preprocessed):
+                print(f"Loading preprocessed few-shot data from {preprocessed}")
+                with open(preprocessed, "rb") as file:
+                    data = pickle.load(file)
+                    train, val = data["train"], data["val"]
+            else:
+                train = self.generate_fewshot_dataset(train, num_shots=num_shots)
+                val = self.generate_fewshot_dataset(val, num_shots=min(num_shots, 4))
+                data = {"train": train, "val": val}
+                print(f"Saving preprocessed few-shot data to {preprocessed}")
+                with open(preprocessed, "wb") as file:
+                    pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)
+
+        # Optionally subsample classes (keep interface consistent)
+        train, val, test = self.subsample_classes(train, val, test, subsample=subsample_classes)
+
+        self.templates = UCF101_TEMPLATES
+
+        super().__init__(train_x=train, val=val, test=test)
+
+    @staticmethod
+    def subsample_classes(*args, subsample="all"):
+        assert subsample in ["all", "base", "new"]
+
+        if subsample == "all":
+            return args
+
+        dataset = args[0]
+        labels = set()
+        for item in dataset:
+            labels.add(item.label)
+        labels = list(labels)
+        labels.sort()
+        n = len(labels)
+        m = math.ceil(n / 2)
+
+        print(f"SUBSAMPLE {subsample.upper()} CLASSES!")
+        if subsample == "base":
+            selected = labels[:m]
+        else:
+            selected = labels[m:]
+        relabeler = {y: y_new for y_new, y in enumerate(selected)}
+
+        output = []
+        for dataset in args:
+            dataset_new = []
+            for item in dataset:
+                if item.label not in selected:
+                    continue
+                item_new = Datum(
+                    impath=item.impath,
+                    label=relabeler[item.label],
+                    classname=item.classname
+                )
+                dataset_new.append(item_new)
+            output.append(dataset_new)
+
+        return output
+
+    # ---- IO helpers (compatible with OxfordPets) ----
+    @staticmethod
+    def save_split(train, val, test, filepath, path_prefix):
+        def _extract(items):
+            out = []
+            for item in items:
+                impath = item.impath
+                label = item.label
+                classname = item.classname
+                impath = impath.replace(path_prefix, "")
+                if impath.startswith("/"):
+                    impath = impath[1:]
+                out.append((impath, label, classname))
+            return out
+
+        train = _extract(train)
+        val = _extract(val)
+        test = _extract(test)
+        split = {"train": train, "val": val, "test": test}
+        write_json(split, filepath)
+        print(f"Saved split to {filepath}")
+
+    @staticmethod
+    def read_split(filepath, path_prefix):
+        def _convert(items):
+            out = []
+            for impath, label, classname in items:
+                impath = os.path.join(path_prefix, impath)
+                item = Datum(impath=impath, label=int(label), classname=classname)
+                out.append(item)
+            return out
+
+        print(f"Reading split from {filepath}")
+        split = read_json(filepath)
+        train = _convert(split["train"])
+        val = _convert(split["val"])
+        test = _convert(split["test"])
+        return train, val, test
+
+    @staticmethod
+    def split_trainval(trainval: List[Datum], p_val=0.2) -> Tuple[List[Datum], List[Datum]]:
+        p_trn = 1 - p_val
+        tracker = defaultdict(list)
+        for idx, item in enumerate(trainval):
+            tracker[item.label].append(idx)
+        train, val = [], []
+        for label, idxs in tracker.items():
+            n_val = max(1, round(len(idxs) * p_val))
+            random.shuffle(idxs)
+            for n, idx in enumerate(idxs):
+                item = trainval[idx]
+                if n < n_val:
+                    val.append(item)
+                else:
+                    train.append(item)
+        return train, val
+
+    # ---- Directory reader (fallback if JSON is missing) ----
+    def _read_from_dir(self) -> List[Datum]:
+        if not os.path.isdir(self.image_dir):
+            raise FileNotFoundError(f"Image directory not found: {self.image_dir}")
+
+        # Build canonical key mapping (lowercased, remove spaces/underscores)
+        def canon_key(s: str) -> str:
+            return ''.join(ch for ch in s.lower() if ch.isalnum())
+
+        canonical_to_label = {name: idx for idx, name in enumerate(UCF101_CANONICAL)}
+        key_to_canonical = {canon_key(name): name for name in UCF101_CANONICAL}
+
+        items: List[Datum] = []
+        class_dirs = listdir_nohidden(self.image_dir, sort=True)
+        for cls_dir in class_dirs:
+            cls_path = os.path.join(self.image_dir, cls_dir)
+            if not os.path.isdir(cls_path):
+                continue
+            # Try to map folder name to canonical class
+            k = canon_key(cls_dir.replace('_', ' '))
+            cname = key_to_canonical.get(k, None)
+            if cname is None:
+                # Try removing underscores without space
+                k2 = canon_key(cls_dir.replace('_', ''))
+                cname = key_to_canonical.get(k2, None)
+            if cname is None:
+                # As a last resort, use the folder name with underscores replaced
+                cname = cls_dir.replace('_', ' ').strip()
+                if cname not in canonical_to_label:
+                    # Unknown class; skip
+                    continue
+            label = canonical_to_label[cname]
+            # Collect images
+            for vid in listdir_nohidden(cls_path, sort=False):
+                vpath = os.path.join(cls_path, vid)
+                if os.path.isdir(vpath):
+                    # some datasets might nest frames under video folder; include all frames
+                    for frame in listdir_nohidden(vpath, sort=False):
+                        impath = os.path.join(vpath, frame)
+                        if os.path.isfile(impath):
+                            items.append(Datum(impath=impath, label=label, classname=cname))
+                else:
+                    # direct frames under class folder
+                    if os.path.isfile(vpath):
+                        items.append(Datum(impath=vpath, label=label, classname=cname))
+        return items

MTIL_datasets/utils.py

@@ -0,0 +1,299 @@
+import os
+import random
+import os.path as osp
+import tarfile
+import zipfile
+from collections import defaultdict
+import gdown
+import errno
+import warnings
+import json
+import numpy as np
+import h5py
+from torch.utils.data import Dataset as TorchDataset
+from PIL import Image
+
+
+class Datum:
+    def __init__(self, impath="", label=0, domain=0, classname=""):
+        self._impath = impath
+        self._label = label
+        self._domain = domain
+        self._classname = classname
+
+    @property
+    def impath(self):
+        return self._impath
+
+    @property
+    def label(self):
+        return self._label
+
+    @property
+    def domain(self):
+        return self._domain
+
+    @property
+    def classname(self):
+        return self._classname
+
+
+class DatasetBase:
+    dataset_dir = "" 
+    domains = []
+
+    def __init__(self, train_x=None, train_u=None, val=None, test=None):
+        self._train_x = train_x
+        self._train_u = train_u
+        self._val = val
+        self._test = test
+        self._num_classes = self.get_num_classes(train_x)
+        self._lab2cname, self._classnames = self.get_lab2cname(train_x)
+
+    @property
+    def train_x(self):
+        return self._train_x
+
+    @property
+    def train_u(self):
+        return self._train_u
+
+    @property
+    def val(self):
+        return self._val
+
+    @property
+    def test(self):
+        return self._test
+
+    @property
+    def lab2cname(self):
+        return self._lab2cname
+
+    @property
+    def classnames(self):
+        return self._classnames
+
+    @property
+    def num_classes(self):
+        return self._num_classes
+
+    @staticmethod
+    def get_num_classes(data_source):
+        label_set = set()
+        for item in data_source:
+            label_set.add(item.label)
+        return max(label_set) + 1
+
+    @staticmethod
+    def get_lab2cname(data_source):
+        container = set()
+        for item in data_source:
+            container.add((item.label, item.classname))
+        mapping = {label: classname for label, classname in container}
+        labels = list(mapping.keys())
+        labels.sort()
+        classnames = [mapping[label] for label in labels]
+        return mapping, classnames
+
+    @property
+    def template(self):
+        return self.templates[0]
+
+    def check_input_domains(self, source_domains, target_domains):
+        assert len(source_domains) > 0, "source_domains (list) is empty"
+        assert len(target_domains) > 0, "target_domains (list) is empty"
+        self.is_input_domain_valid(source_domains)
+        self.is_input_domain_valid(target_domains)
+
+    def is_input_domain_valid(self, input_domains):
+        for domain in input_domains:
+            if domain not in self.domains:
+                raise ValueError(
+                    "Input domain must belong to {}, "
+                    "but got [{}]".format(self.domains, domain)
+                )
+
+    def download_data(self, url, dst, from_gdrive=True):
+        if not osp.exists(osp.dirname(dst)):
+            os.makedirs(osp.dirname(dst))
+
+        if from_gdrive:
+            gdown.download(url, dst, quiet=False)
+        else:
+            raise NotImplementedError
+
+        print("Extracting file ...")
+
+        if dst.endswith(".zip"):
+            zip_ref = zipfile.ZipFile(dst, "r")
+            zip_ref.extractall(osp.dirname(dst))
+            zip_ref.close()
+
+        elif dst.endswith(".tar"):
+            tar = tarfile.open(dst, "r:")
+            tar.extractall(osp.dirname(dst))
+            tar.close()
+
+        elif dst.endswith(".tar.gz"):
+            tar = tarfile.open(dst, "r:gz")
+            tar.extractall(osp.dirname(dst))
+            tar.close()
+
+        else:
+            raise NotImplementedError
+
+        print("File extracted to {}".format(osp.dirname(dst)))
+
+    def generate_fewshot_dataset(
+        self, *data_sources, num_shots=-1, repeat=False
+    ):
+        if num_shots < 1:
+            if len(data_sources) == 1:
+                return data_sources[0]
+            return data_sources
+
+        print(f"Creating a {num_shots}-shot dataset")
+
+        output = []
+
+        for data_source in data_sources:
+            tracker = self.split_dataset_by_label(data_source)
+            dataset = []
+
+            for label, items in tracker.items():
+                if len(items) >= num_shots:
+                    sampled_items = random.sample(items, num_shots)
+                else:
+                    if repeat:
+                        sampled_items = random.choices(items, k=num_shots)
+                    else:
+                        sampled_items = items
+                dataset.extend(sampled_items)
+
+            output.append(dataset)
+
+        if len(output) == 1:
+            return output[0]
+
+        return output
+
+    def split_dataset_by_label(self, data_source):
+        output = defaultdict(list)
+
+        for item in data_source:
+            output[item.label].append(item)
+
+        return output
+
+    def split_dataset_by_domain(self, data_source):
+        output = defaultdict(list)
+
+        for item in data_source:
+            output[item.domain].append(item)
+
+        return output
+
+
+class DatasetWrapper(TorchDataset):
+
+    def __init__(self, data_source, transform=None, is_train=False):
+        self.data_source = data_source
+        self.transform = transform
+        self.is_train = is_train
+        self._h5_cache = {}
+        self._h5_info_printed = set()
+
+    def __len__(self):
+        return len(self.data_source)
+
+    def __getitem__(self, idx):
+        item = self.data_source[idx]
+
+        impath = item.impath
+        if isinstance(impath, str):
+            img0 = Image.open(impath).convert("RGB")
+        elif isinstance(impath, tuple) and len(impath) == 4 and impath[0] == 'h5':
+            _, fpath, key, index = impath
+            f = self._h5_cache.get(fpath)
+            if f is None:
+                f = h5py.File(fpath, 'r')
+                self._h5_cache[fpath] = f
+                # Print file info once
+                try:
+                    keys = list(f.keys())
+                    print(f"H5 open: {os.path.basename(fpath)} keys={keys[:5]}{'...' if len(keys) > 5 else ''}")
+                except Exception as e:
+                    print(f"H5 open (keys) failed for {fpath}: {e}")
+            # Print dataset info per file the first time we see this path
+            if fpath not in self._h5_info_printed:
+                try:
+                    ds = f[key]
+                    print(f"H5 dataset: {os.path.basename(fpath)}[{key}] shape={getattr(ds, 'shape', '?')} dtype={getattr(ds, 'dtype', '?')}")
+                except Exception as e:
+                    print(f"H5 dataset info failed for {fpath}[{key}]: {e}")
+                self._h5_info_printed.add(fpath)
+            arr = f[key][int(index)]
+            arr = np.asarray(arr)
+            # Convert CHW -> HWC if needed
+            if arr.ndim == 3 and arr.shape[0] in (1, 3) and arr.shape[-1] not in (1, 3):
+                arr = np.transpose(arr, (1, 2, 0))
+            # Ensure HWC and uint8
+            if arr.ndim == 3 and arr.shape[-1] in (1, 3):
+                pass
+            else:
+                raise ValueError(f"Unexpected H5 image shape: {arr.shape}")
+            if arr.dtype != np.uint8:
+                arr = arr.astype(np.uint8)
+            if arr.shape[-1] == 1:
+                img0 = Image.fromarray(arr.squeeze(-1), mode='L').convert('RGB')
+            else:
+                img0 = Image.fromarray(arr, mode='RGB')
+        else:
+            # Fallback: if already PIL Image
+            if isinstance(impath, Image.Image):
+                img0 = impath
+            else:
+                raise ValueError("Unsupported impath type in DatasetWrapper")
+
+        if self.transform:
+            img = self.transform(img0)
+        else:
+            img = img0
+
+        return img, item.label
+
+
+def check_isfile(fpath):
+    isfile = osp.isfile(fpath)
+    if not isfile:
+        warnings.warn('No file found at "{}"'.format(fpath))
+    return isfile
+
+
+def mkdir_if_missing(dirname):
+    if not osp.exists(dirname):
+        try:
+            os.makedirs(dirname)
+        except OSError as e:
+            if e.errno != errno.EEXIST:
+                raise
+
+
+def listdir_nohidden(path, sort=False):
+    items = [f for f in os.listdir(path) if not f.startswith(".")]
+    if sort:
+        items.sort()
+    return items
+
+
+def read_json(fpath):
+    with open(fpath, "r") as f:
+        obj = json.load(f)
+    return obj
+
+
+def write_json(obj, fpath):
+    mkdir_if_missing(osp.dirname(fpath))
+    with open(fpath, "w") as f:
+        json.dump(obj, f, indent=4, separators=(",", ": "))

MTIL_datasets/voc2007.py

@@ -0,0 +1,173 @@
+import os
+import os.path as osp
+from typing import List, Dict, Tuple
+from collections import defaultdict
+
+from .utils import Datum, DatasetBase, listdir_nohidden
+
+
+VOC2007_CLASSES = [
+    'aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
+    'bus', 'car', 'cat', 'chair', 'cow',
+    'dog', 'horse', 'motorbike', 'person', 'sheep',
+    'sofa', 'diningtable', 'pottedplant', 'train', 'tvmonitor',
+]
+
+VOC2007_TEMPLATES = [
+    'a photo of a {}.',
+]
+
+
+class VOC2007(DatasetBase):
+    """
+    VOC2007 multi-label classification dataset adapter for MTIL evaluation.
+
+    Expects directory structure:
+      data/VOC2007/
+        JPEGImages/
+        Main/
+          <class>_trainval.txt
+          <class>_test.txt
+
+    Each *_split.txt has lines: "<image_id> <label>", where label in {1, 0, -1}.
+    We treat 1 as positive and others as negative for that class.
+    """
+
+    dataset_dir = 'VOC2007'
+
+    def __init__(self, root: str, seed: int = 32, single_label: bool = False):
+        self.root = os.path.expanduser(root)
+        self.dataset_dir = osp.join(self.root, self.dataset_dir)
+        self.images_dir = osp.join(self.dataset_dir, 'JPEGImages')
+        self.main_dir = osp.join(self.dataset_dir, 'Main')
+
+        classes = VOC2007_CLASSES
+        class_to_label = {c: i for i, c in enumerate(classes)}
+
+        if not single_label:
+            # Multi-label mode (for Zero-Shot evaluation)
+            test_list = self._read_split('test', classes)
+            test = self._build_data_list(test_list, class_to_label)
+
+            # Create a meta train set with one sample per class (integer labels) for proper initialization
+            # Try to point to a real image containing that class as positive; fallback to any available image or empty path
+            meta_train: List[Datum] = []
+            for ci, cname in enumerate(classes):
+                impath = None
+                for d in test:
+                    try:
+                        if isinstance(d.label, (list, tuple)) and ci < len(d.label) and int(d.label[ci]) == 1:
+                            impath = d.impath
+                            break
+                    except Exception:
+                        continue
+                if impath is None:
+                    impath = test[0].impath if len(test) > 0 else ''
+                meta_train.append(Datum(impath=impath, label=ci, classname=cname))
+
+            train_x: List[Datum] = meta_train
+            val: List[Datum] = []
+            self.templates = VOC2007_TEMPLATES
+            super().__init__(train_x=train_x, val=val, test=test)
+        else:
+            # Single-label mode (for CIL downstream training/evaluation)
+            # Build trainval and test splits as single-label datasets by selecting the first positive class per image
+            trainval_multi = self._read_split('trainval', classes)
+            test_multi = self._read_split('test', classes)
+            train_x = self._build_single_label_list(trainval_multi, classes)
+            val = []
+            test = self._build_single_label_list(test_multi, classes)
+            self.templates = VOC2007_TEMPLATES
+            super().__init__(train_x=train_x, val=val, test=test)
+
+    def _read_split(self, split: str, classes: List[str]) -> Dict[str, List[int]]:
+        """Return mapping: image_id -> multi-hot list for the given split."""
+        # Collect image ids present in this split
+        img_ids = set()
+        by_class_labels: Dict[str, Dict[str, int]] = {}
+        for cname in classes:
+            split_file = osp.join(self.main_dir, f"{cname}_{split}.txt")
+            if not osp.isfile(split_file):
+                # Try alternate common path name
+                split_file = osp.join(self.dataset_dir, 'ImageSets', 'Main', f"{cname}_{split}.txt")
+            if not osp.isfile(split_file):
+                # If missing, skip this class
+                continue
+            class_map: Dict[str, int] = {}
+            with open(split_file, 'r') as f:
+                for line in f:
+                    parts = line.strip().split()
+                    if len(parts) < 2:
+                        continue
+                    img_id, label_str = parts[0], parts[1]
+                    try:
+                        label = int(label_str)
+                    except Exception:
+                        continue
+                    class_map[img_id] = 1 if label == 1 else 0
+                    img_ids.add(img_id)
+            by_class_labels[cname] = class_map
+
+        # Aggregate into multi-hot vectors
+        result: Dict[str, List[int]] = {}
+        for img_id in img_ids:
+            vec = [0] * len(classes)
+            for ci, cname in enumerate(classes):
+                cmap = by_class_labels.get(cname, {})
+                vec[ci] = int(cmap.get(img_id, 0))
+            result[img_id] = vec
+        return result
+
+    def _build_data_list(self, id_to_vec: Dict[str, List[int]], class_to_label: Dict[str, int]) -> List[Datum]:
+        data: List[Datum] = []
+        for img_id, vec in id_to_vec.items():
+            # Try common image extensions
+            impath = None
+            for ext in ['.jpg', '.jpeg', '.png']:
+                p = osp.join(self.images_dir, img_id + ext)
+                if osp.isfile(p):
+                    impath = p
+                    break
+            if impath is None:
+                # As a fallback, if there's exactly one file starting with img_id
+                try:
+                    candidates = [f for f in listdir_nohidden(self.images_dir) if f.startswith(img_id + '.')]
+                    if candidates:
+                        impath = osp.join(self.images_dir, candidates[0])
+                except Exception:
+                    pass
+            if impath is None:
+                # Skip missing images
+                continue
+            # For multi-label, we store the full vector as the label
+            data.append(Datum(impath=impath, label=vec, classname=''))
+        return data
+
+    def _build_single_label_list(self, id_to_vec: Dict[str, List[int]], classes: List[str]) -> List[Datum]:
+        data: List[Datum] = []
+        for img_id, vec in id_to_vec.items():
+            # choose the first positive class; skip if none
+            try:
+                cls_idx = next((i for i, v in enumerate(vec) if int(v) == 1), None)
+            except Exception:
+                cls_idx = None
+            if cls_idx is None:
+                continue
+            # image path resolution
+            impath = None
+            for ext in ['.jpg', '.jpeg', '.png']:
+                p = osp.join(self.images_dir, img_id + ext)
+                if osp.isfile(p):
+                    impath = p
+                    break
+            if impath is None:
+                try:
+                    candidates = [f for f in listdir_nohidden(self.images_dir) if f.startswith(img_id + '.')]
+                    if candidates:
+                        impath = osp.join(self.images_dir, candidates[0])
+                except Exception:
+                    pass
+            if impath is None:
+                continue
+            data.append(Datum(impath=impath, label=int(cls_idx), classname=classes[cls_idx]))
+        return data

README.md

@@ -0,0 +1,227 @@
+# DFA-MoE: Tackling Dual Forgetting in Vision-Language Continual Learning
+
+[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/)
+[![PyTorch](https://img.shields.io/badge/PyTorch-red.svg)](https://pytorch.org/)
+
+Official PyTorch implementation of the paper "Don't Forget Why You Started: Tackling Dual Forgetting in Vision-Language Continual Learning" ICML2026.
+
+## Abstract
+
+Vision-Language Models (VLMs) are strong continual learners, but standard class-incremental learning often damages the pre-trained vision-language alignment that underpins zero-shot generalization. This repository studies that dual-forgetting problem from two perspectives: Incremental Knowledge Forgetting (IKF), which harms previously learned classes, and Pre-trained Knowledge Forgetting (PKF), which erodes the original zero-shot capabilities of the VLM. To address this issue, we implement the Dual-Forgetting-Aware Class-Incremental Learning (DFA-CIL) framework, the Similarity-Calibrated Retention (SCR) metric, and DFA-MoE, a functionally heterogeneous PEFT method that decouples alignment preservation from task adaptation.
+
+![Introduction](docs/intro.png)
+
+## Key Contributions
+
+- Dual-forgetting evaluation for VLM continual learning through the DFA-CIL protocol.
+- Similarity-Calibrated Retention (SCR) utilizing similarity based weight to disentangle genuine foundational retention from the confounding effects of positive transfer.
+- DFA-MoE with two functional pathways:
+  - Alignment Pathway: task-agnostic contrastive expert for PKF mitigation.
+  - Plasticity Pathway: task-specific experts with classification and auxiliary contrastive learning for IKF mitigation.
+  - Hierarchical routing with an inner router over task-specific experts and an outer router balancing alignment and plasticity outputs.
+
+## Installation
+
+### Setup Environment
+
+```bash
+git clone https://github.com/RL-MIND/DFA-MoE
+cd DFA_MoE
+pip install -r requirements.txt
+```
+
+## Supported Downstream Datasets
+
+The current `configs/class` directory provides downstream class-incremental configs for the following dataset indices. These indices are the values used in `+train_dataset=[id]`.
+
+- `0`: `FGVCAircraft`
+- `1`: `Caltech101`
+- `2`: `CIFAR100`
+- `3`: `DescribableTextures`
+- `4`: `EuroSAT`
+- `5`: `OxfordFlowers`
+- `6`: `Food101`
+- `7`: `MNIST`
+- `8`: `OxfordPets`
+- `9`: `StanfordCars`
+- `10`: `SUN397`
+- `11`: `Country211`
+- `14`: `GTSRB`
+- `15`: `RESISC45`
+- `16`: `FER2013`
+- `17`: `UCF101`
+- `18`: `CIFAR10`
+- `19`: `STL10`
+- `20`: `VOC2007`
+- `21`: `ImageNetR`
+- `22`: `KittiDistance`
+- `24`: `CLEVRCount`
+
+The datasets pool also contains:
+
+- `12`: `SST2`
+- `13`: `HatefulMemes`
+- `23`: `PCam`
+
+These datasets are not supported as `--downstream-dataset` for `calculate_sim.py` because they are binary classification datasets.
+
+## Data Preparation
+
+We will soon release all the datasets used in this work at [google drive](https://drive.google.com/drive/folders/1DxD6MixpyTcsLSv8Gah5V0Cp3yMupYZO?usp=sharing).
+
+1. Download or prepare the datasets under a common `dataset_root`.
+2. Update `dataset_root` in the command line when launching experiments.
+3. Choose the dataset-specific config from `configs/class/`.
+
+Example dataset config files:
+
+- `configs/class/eurosat.yaml`
+- `configs/class/flower.yaml`
+- `configs/class/cifar100.yaml`
+- `configs/class/aircraft.yaml`
+
+## Quick Start
+
+### Train a DFA-CIL Run
+
+Example: EuroSAT with 5 class-incremental splits.
+
+```bash
+python main.py \
+  --config-path ./configs/class \
+  --config-name eurosat.yaml \
+  dataset_root="/path/to/data" \
+  +train_dataset=[4] \
+  +cil_splits=[5]
+```
+
+Example: OxfordFlowers with 17 splits.
+
+```bash
+python main.py \
+  --config-path ./configs/class \
+  --config-name flower.yaml \
+  dataset_root="/path/to/data" \
+  +train_dataset=[5] \
+  +cil_splits=[17]
+```
+
+
+## Important Configuration Options
+
+This project uses Hydra-based configuration. Key parameters include:
+
+```yaml
+model_name: "ViT-B/16"
+prompt_template: "a bad photo of a {}."
+
+batch_size: 128
+weight_decay: 0.0
+ls: 0.0
+
+epochs_a: 1
+epochs_b: 1
+
+lr_e1: 1.0e-5
+lr_e2: 5.0e-4
+text_lr_e2: 1.0e-5
+lr_e2_router: 1.0e-3
+lr_top_router: 1.0e-5
+
+tau_con: 0.15
+tau_b_con: 0.17
+lambda_b_con: 0.001
+
+num_task_experts: 2
+e2_top_k: 2
+moco_queue_size: 128
+```
+
+## Pre-task Zero-shot Baseline
+
+To compute SCR correctly, the metric log must contain the original CLIP zero-shot baseline `A_k^0`. Enable this by setting:
+
+```yaml
+pre_task_zero_shot_eval: true
+zero_shot_eval: true
+```
+
+or from the command line:
+
+```bash
+python main.py \
+  --config-path ./configs/class \
+  --config-name eurosat.yaml \
+  dataset_root="/path/to/data" \
+  +train_dataset=[4] \
+  +cil_splits=[5] \
+  pre_task_zero_shot_eval=true \
+  zero_shot_eval=true
+```
+
+This writes a `task: -1` entry with `zs_pre` into `metrics.json`, which is required by `calculate_SCR.py`.
+
+## SCR Evaluation Workflow
+
+### Step 1: Run continual learning and save `metrics.json`
+
+Run `main.py` with `pre_task_zero_shot_eval=true` and `zero_shot_eval=true`.
+
+### Step 2: Generate the similarity matrix
+
+Example: EuroSAT with 5 splits.
+
+```bash
+python calculate_sim.py \
+  --dataset-root "/path/to/data" \
+  --downstream-dataset EuroSAT \
+  --cil-split 5 \
+  --output eurosat_similarity.json
+```
+
+This script:
+
+- uses the original frozen CLIP model,
+- produces a task-to-upstream similarity matrix for SCR.
+
+### Step 3: Compute SCR
+
+```bash
+python calculate_SCR.py \
+  --metric-path /path/to/metrics.json \
+  --sim-json /path/to/eurosat_similarity_sim.json
+```
+
+`calculate_SCR.py` now strictly requires:
+
+- a `task: -1` row with non-empty `zs_pre`,
+- per-task `zs` results in `metrics.json`,
+- a similarity JSON generated by `calculate_sim.py`.
+
+## Reference Script
+
+The repository also contains `run.sh` as a reference batch script.
+
+## Notes
+
+- `main.py`, `calculate_sim.py`, and `calculate_SCR.py` are the recommended entry points for experiments and evaluation.
+
+## Citation
+
+If you find this repository useful in your research, please cite the paper:
+
+```bibtex
+@inproceedings{kang2026dont,
+  title={Don't Forget Why You Started: Tackling Dual Forgetting in Vision-Language Continual Learning},
+  author={Kang, Borui and Gu, Jinrui and Feng, Tao and Fan, Qi and Shi, Yinghuan and Wang, Lei and Li, Wenbin and Gao, Yang},
+  booktitle={Proceedings of the 43rd International Conference on Machine Learning},
+  year={2026}
+}
+```
+
+## Acknowledgement
+
+This repository is built upon and modified from [MoE-Adapters4CL](https://github.com/JiazuoYu/MoE-Adapters4CL). We thank the original authors for making their code publicly available.
+
+The dataset processing pipeline in this repository is also based on [DIKI](https://github.com/lloongx/DIKI). We thank the authors for releasing their implementation.
+

calculate_SCR.py

@@ -0,0 +1,266 @@
+from __future__ import annotations
+
+import argparse
+import json
+from typing import List, Dict, Any, Tuple
+
+
+DATASET_NAME_ALIASES = {
+    "Aircraft": "FGVCAircraft",
+    "FGVCAircraft": "FGVCAircraft",
+    "DTD": "DescribableTextures",
+    "DescribableTextures": "DescribableTextures",
+}
+
+
+def canonicalize_dataset_name(name: str) -> str:
+    return DATASET_NAME_ALIASES.get(str(name), str(name))
+
+
+def format_matrix(mat: List[List[float]], decimals: int = 3) -> str:
+    if not mat:
+        return "[]"
+    fmt = f"{{:.{decimals}f}}"
+    lines = []
+    for row in mat:
+        row_str = ", ".join(fmt.format(v) for v in row)
+        lines.append(f"[{row_str}]")
+    return "[" + "\n" + ",\n".join(lines) + "\n]"
+
+
+def load_lines(path: str) -> List[Dict[str, Any]]:
+    rows = []
+    with open(path, 'r', encoding='utf-8') as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                rows.append(json.loads(line))
+            except Exception as e:
+                print(f"[WARN] skip invalid JSON line: {e}")
+    rows.sort(key=lambda r: int(r.get('task', 0)))
+    return rows
+
+
+def build_zs_matrix(rows: List[Dict[str, Any]]) -> Tuple[List[List[float]], List[str]]:
+    if not rows:
+        return [], []
+
+    baseline_row = None
+    task_rows = []
+
+    for r in rows:
+        task_id = int(r.get('task', 0))
+        zs_pre = r.get('zs_pre')
+        zs_task = r.get('zs')
+        zs_task_alt = r.get('zs_mtil')
+
+        if task_id == -1 and isinstance(zs_pre, dict) and len(zs_pre) > 0:
+            baseline_row = zs_pre
+            continue
+        if task_id >= 0:
+            if isinstance(zs_task, dict) and len(zs_task) > 0:
+                task_rows.append((task_id, zs_task))
+            elif isinstance(zs_task_alt, dict) and len(zs_task_alt) > 0:
+                task_rows.append((task_id, zs_task_alt))
+
+    if baseline_row is None:
+        raise ValueError(
+            "metric.json must contain a pre-task baseline row with `task: -1` and non-empty `zs_pre`."
+        )
+    if not task_rows:
+        return [], []
+
+    names = list(baseline_row.keys())
+    matrix: List[List[float]] = []
+    matrix.append([float(baseline_row.get(name, 0.0)) for name in names])
+    for _task_id, zs in sorted(task_rows, key=lambda item: item[0]):
+        row = [float(zs.get(name, 0.0)) for name in names]
+        matrix.append(row)
+    return matrix, names
+
+
+def load_similarity_json(path: str) -> Tuple[List[List[float]], List[str]]:
+    with open(path, 'r', encoding='utf-8') as f:
+        payload = json.load(f)
+    sim_mat_raw = payload.get("similarity_matrix", [])
+    sim_names_raw = payload.get("upstream_names", [])
+    sim_mat = [[float(x) for x in row] for row in sim_mat_raw]
+    sim_names = [canonicalize_dataset_name(x) for x in sim_names_raw]
+    return sim_mat, sim_names
+
+
+def compute_scores_with_matrix(
+    zs_mat: List[List[float]],
+    zs_names: List[str],
+    sim_mat: List[List[float]],
+    sim_names: List[str],
+) -> List[float]:
+    if not zs_mat or not sim_mat:
+        return []
+
+    zs_names = [canonicalize_dataset_name(nm) for nm in zs_names]
+    zs_name_to_idx = {nm: i for i, nm in enumerate(zs_names)}
+    sim_name_to_idx = {nm: i for i, nm in enumerate(sim_names)}
+    common = [nm for nm in zs_names if nm in sim_name_to_idx]
+    if len(common) < len(zs_names):
+        missing = [nm for nm in zs_names if nm not in sim_name_to_idx]
+        print(f"[WARN] Missing similarity columns for: {missing}")
+    zs_cols = [zs_name_to_idx[nm] for nm in common]
+    sim_cols = [sim_name_to_idx[nm] for nm in common]
+
+    base = zs_mat[0]
+    scores: List[float] = []
+    T = min(len(sim_mat), max(0, len(zs_mat) - 1))
+    if T < (len(zs_mat) - 1):
+        print(f"[WARN] sim-matrix rows ({len(sim_mat)}) < tasks ({len(zs_mat) - 1}); truncating to {T}")
+    for t in range(T):
+        row = zs_mat[t + 1]
+        total = 0.0
+        sum_w = 0.0
+        sims = [float(sim_mat[t][sim_cols[k]]) for k in range(len(common))]
+        risks = [max(0.0, 1-s) for s in sims]
+        norm = sum(risks) if risks else 0.0
+        default_weight = 1.0 / max(1, len(common))
+        for k in range(len(common)):
+            i = zs_cols[k]
+            diff = float(row[i]) - float(base[i])
+            w = (risks[k] / norm) if norm > 0.0 else default_weight
+            total += diff * w
+            sum_w += w
+        avg = (total / sum_w) if sum_w > 0.0 else 0.0
+        scores.append(round(avg, 2))
+    return scores
+
+
+def compute_scores_grouped(
+    zs_mat: List[List[float]],
+    zs_names: List[str],
+    sim_mat: List[List[float]],
+    sim_names: List[str],
+) -> Dict[str, Any]:
+    if not zs_mat or not sim_mat:
+        return {
+            "weighted": {"low": [], "mid": [], "high": []},
+        }
+
+    zs_names = [canonicalize_dataset_name(nm) for nm in zs_names]
+    zs_name_to_idx = {nm: i for i, nm in enumerate(zs_names)}
+    sim_name_to_idx = {nm: i for i, nm in enumerate(sim_names)}
+    common = [nm for nm in zs_names if nm in sim_name_to_idx]
+    if len(common) < len(zs_names):
+        missing = [nm for nm in zs_names if nm not in sim_name_to_idx]
+        print(f"[WARN] Missing similarity columns for grouped score: {missing}")
+    zs_cols = [zs_name_to_idx[nm] for nm in common]
+    sim_cols = [sim_name_to_idx[nm] for nm in common]
+
+    base = zs_mat[0]
+    T = min(len(sim_mat), max(0, len(zs_mat) - 1))
+    if T < (len(zs_mat) - 1):
+        print(f"[WARN] sim-matrix rows ({len(sim_mat)}) < tasks ({len(zs_mat) - 1}); truncating to {T}")
+
+    def calc_weighted(items):
+        total = sum(item["diff"] * item["weight"] for item in items)
+        sum_w = sum(item["weight"] for item in items)
+        return round((total / sum_w) if sum_w > 0.0 else 0.0, 2)
+
+    w_low, w_mid, w_high = [], [], []
+
+    for t in range(T):
+        row = zs_mat[t + 1]
+        sims = [float(sim_mat[t][sim_cols[k]]) for k in range(len(common))]
+        risks = [max(0.0, 1.0 - s) for s in sims]
+        norm = sum(risks) if risks else 0.0
+        default_weight = 1.0 / max(1, len(common))
+        triplets = []
+        for k in range(len(common)):
+            i = zs_cols[k]
+            sim_val = sims[k]
+            diff = float(row[i]) - float(base[i])
+            w = (risks[k] / norm) if norm > 0.0 else default_weight
+            triplets.append({
+                "sim": sim_val,
+                "diff": diff,
+                "weight": w,
+                "index": i + 1,
+            })
+        triplets.sort(key=lambda x: x["sim"])
+        n = len(triplets)
+        q = n // 3
+        low = triplets[:q]
+        mid = triplets[q:2 * q]
+        high = triplets[2 * q:]
+
+        w_low.append(calc_weighted(low))
+        w_mid.append(calc_weighted(mid))
+        w_high.append(calc_weighted(high))
+
+    return {
+        "weighted": {"low": w_low, "mid": w_mid, "high": w_high},
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Compute SCR metrics from metric.json and calculate_sim similarity JSON.",
+    )
+    parser.add_argument("--metric-path", type=str, default="metric.json", help="Path to the metric.json file")
+    parser.add_argument("--sim-json", type=str, required=True, help="Path to the similarity JSON generated by calculate_sim.py")
+    parser.add_argument("--zs-decimals", type=int, default=2, help="Decimals when printing ZS matrix")
+    args = parser.parse_args()
+
+    rows = load_lines(args.metric_path)
+    try:
+        zs_mat, zs_names = build_zs_matrix(rows)
+    except ValueError as e:
+        print(f"[ERROR] {e}")
+        return
+    sim_mat, sim_names = load_similarity_json(args.sim_json)
+
+    if not zs_mat:
+        print("[ERROR] No zero-shot matrix could be extracted from metric.json.")
+        return
+    if not sim_mat:
+        print("[ERROR] No similarity matrix could be extracted from the calculate_sim JSON.")
+        return
+
+    num_cols = len(zs_mat[0]) if zs_mat else 0
+    if any(len(row) != num_cols for row in zs_mat):
+        print("[ERROR] Extracted zero-shot matrix rows must all have the same length.")
+        return
+
+    sim_cols = len(sim_mat[0]) if sim_mat else 0
+    if any(len(row) != sim_cols for row in sim_mat):
+        print("[ERROR] similarity_matrix rows in the calculate_sim JSON must all have the same length.")
+        return
+    if sim_names and len(sim_names) != sim_cols:
+        print("[ERROR] upstream_names length must match similarity_matrix columns.")
+        return
+    if not sim_names:
+        sim_names = zs_names
+
+    print("1. Zero-shot matrix (baseline + tasks):")
+    print(format_matrix(zs_mat, decimals=args.zs_decimals))
+
+    scores_weighted = compute_scores_with_matrix(zs_mat, zs_names, sim_mat, sim_names)
+    avg_weighted = round(sum(scores_weighted) / len(scores_weighted), 2) if scores_weighted else 0.0
+    print("\nSCR:")
+    print(avg_weighted)
+
+    grouped = compute_scores_grouped(zs_mat, zs_names, sim_mat, sim_names)
+
+    def mean_or_zero(arr: list) -> float:
+        return round(sum(arr) / len(arr), 2) if arr else 0.0
+
+    weighted = grouped["weighted"]
+    low_w_mean = mean_or_zero(weighted["low"])
+    mid_w_mean = mean_or_zero(weighted["mid"])
+    high_w_mean = mean_or_zero(weighted["high"])
+
+    print("\nLow-Similarity SCR:", low_w_mean)
+    print("Mid-Similarity SCR:", mid_w_mean)
+    print("High-Similarity SCR:", high_w_mean)
+
+if __name__ == "__main__":
+    main()

calculate_sim.py

@@ -0,0 +1,456 @@
+import argparse
+import os
+from typing import List, Dict, Tuple
+
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+
+# Project-local imports
+from mtil_datasets import get_dataset as get_mtil_dataset
+from continual_clip.mtil_cil import build_mtil_cil_scenarios
+from continual_clip.clip_original import load as load_orig_clip, tokenize as tokenize_orig
+
+
+DEFAULT_NAMES_ORDER1 = [
+    "FGVCAircraft", "Caltech101", "CIFAR100", "DescribableTextures", "EuroSAT",
+    "OxfordFlowers", "Food101", "MNIST", "OxfordPets", "StanfordCars",
+    "SUN397", "Country211", "SST2", "HatefulMemes", "GTSRB",
+    "RESISC45", "FER2013", "UCF101", "CIFAR10", "STL10",
+    "VOC2007", "ImageNetR", "KittiDistance", "PCam", "CLEVRCount",
+]
+
+
+DATASET_NAME_TO_CONFIG_STEM = {
+    "FGVCAircraft": "aircraft",
+    "Aircraft": "aircraft",
+    "Caltech101": "caltech",
+    "CIFAR100": "cifar100",
+    "DescribableTextures": "dtd",
+    "DTD": "dtd",
+    "EuroSAT": "eurosat",
+    "OxfordFlowers": "flower",
+    "Food101": "food",
+    "MNIST": "mnist",
+    "OxfordPets": "pet",
+    "StanfordCars": "car",
+    "SUN397": "sun",
+    "Country211": "conuntry",
+    "GTSRB": "gtsrb",
+    "RESISC45": "resisc45",
+    "FER2013": "fer2013",
+    "UCF101": "ucf",
+    "CIFAR10": "cifar10",
+    "STL10": "stl",
+    "VOC2007": "voc",
+    "ImageNetR": "image-r",
+    "KittiDistance": "kitti",
+    "CLEVRCount": "clevr",
+}
+
+
+UNSUPPORTED_DOWNSTREAM_DATASETS = {
+    "SST2",
+    "HatefulMemes",
+    "PCam",
+}
+
+
+def resolve_downstream_seed(repo_root: str, downstream_name: str) -> int:
+    if downstream_name in UNSUPPORTED_DOWNSTREAM_DATASETS:
+        raise ValueError(
+            f"Downstream dataset '{downstream_name}' is currently unsupported as `--downstream-dataset` because no matching `configs/class/*.yaml` file is available."
+        )
+    config_stem = DATASET_NAME_TO_CONFIG_STEM.get(downstream_name)
+    if not config_stem:
+        raise ValueError(
+            f"No `configs/class/*.yaml` mapping is defined for downstream dataset '{downstream_name}'."
+        )
+    config_path = os.path.join(repo_root, "configs", "class", f"{config_stem}.yaml")
+    if not os.path.exists(config_path):
+        raise FileNotFoundError(
+            f"Expected downstream config file '{config_path}' for dataset '{downstream_name}'."
+        )
+    cfg = OmegaConf.load(config_path)
+    seed = OmegaConf.select(cfg, "seed")
+    if seed is None:
+        raise ValueError(
+            f"Config file '{config_path}' does not define a `seed` entry."
+        )
+    return int(seed)
+
+
+def encode_texts(model, device, sentences: List[str], batch_size: int = 256) -> torch.Tensor:
+    feats_all = []
+    with torch.no_grad():
+        for i in range(0, len(sentences), batch_size):
+            chunk = sentences[i:i+batch_size]
+            tokens = tokenize_orig(chunk).to(device)
+            feats = model.encode_text(tokens)
+            feats = feats / (feats.norm(dim=-1, keepdim=True) + 1e-12)
+            feats_all.append(feats)
+    if not feats_all:
+        return torch.zeros((0, model.text_projection.shape[1]), device=device)
+    return torch.cat(feats_all, dim=0)
+
+
+def build_text_class_prototypes(model, device, classnames: List[str], templates, batch_size: int = 256) -> torch.Tensor:
+    D = model.text_projection.shape[1]
+    class_vecs: List[torch.Tensor] = []
+    for cname in classnames:
+        sents: List[str] = []
+        if templates and len(templates) > 0:
+            for t in templates:
+                try:
+                    sents.append(t(cname) if callable(t) else str(t).format(cname))
+                except Exception:
+                    continue
+        else:
+            sents = [f"a photo of a {cname}."]
+        feats = encode_texts(model, device, sents, batch_size=batch_size)
+        if feats.numel() == 0:
+            v = torch.zeros(D, device=device)
+        else:
+            v = feats.mean(dim=0)
+            v = v / (v.norm() + 1e-12)
+        class_vecs.append(v)
+    if not class_vecs:
+        return torch.zeros((0, D), device=device)
+    return torch.stack(class_vecs, dim=0)
+
+
+def normalize_labels_to_list(lab) -> List[int]:
+    if isinstance(lab, (int, np.integer)):
+        return [int(lab)]
+    import torch as _torch
+    if _torch.is_tensor(lab):
+        arr = lab.detach().cpu().numpy()
+        if arr.ndim == 0:
+            return [int(arr)]
+        if arr.ndim == 1 and arr.size > 1 and set(np.unique(arr)).issubset({0,1}):
+            return [int(x) for x in np.where(arr > 0.5)[0].tolist()]
+        return [int(x) for x in arr.flatten().tolist()]
+    if isinstance(lab, (list, tuple, np.ndarray)):
+        arr = np.asarray(lab)
+        if arr.ndim == 0:
+            return [int(arr)]
+        if arr.ndim == 1 and arr.size > 1 and set(np.unique(arr)).issubset({0,1}):
+            return [int(x) for x in np.where(arr > 0.5)[0].tolist()]
+        if arr.ndim == 1:
+            return [int(x) for x in arr.tolist()]
+        return [int(x) for x in np.where(arr.flatten() > 0.5)[0].tolist()]
+    return []
+
+
+def build_visual_class_prototypes(model, device, ds_wrapper, num_classes: int, max_per_class: int, batch_size: int, preprocess_eval) -> torch.Tensor:
+    D = None
+    sums: Dict[int, torch.Tensor] = {}
+    counts: Dict[int, int] = {i: 0 for i in range(num_classes)}
+    imgs_batch: List[torch.Tensor] = []
+    labels_batch_multi: List[List[int]] = []
+    with torch.no_grad():
+        for i in range(len(ds_wrapper)):
+            img, lab = ds_wrapper[i]
+            lab_ids = [lid for lid in normalize_labels_to_list(lab) if 0 <= lid < num_classes and counts[lid] < max_per_class]
+            if not lab_ids:
+                continue
+            if isinstance(img, torch.Tensor):
+                tensor_img = img
+            else:
+                tensor_img = preprocess_eval(img)
+            imgs_batch.append(tensor_img.unsqueeze(0))
+            labels_batch_multi.append(lab_ids)
+            if len(imgs_batch) >= max(1, batch_size):
+                batch = torch.cat(imgs_batch, dim=0).to(device)
+                feats = model.encode_image(batch)
+                feats = feats / (feats.norm(dim=-1, keepdim=True) + 1e-12)
+                if D is None:
+                    D = int(feats.shape[1])
+                for f, ls in zip(feats, labels_batch_multi):
+                    for l in ls:
+                        if counts[l] >= max_per_class:
+                            continue
+                        if l not in sums:
+                            sums[l] = f.detach().clone()
+                        else:
+                            sums[l] = sums[l] + f.detach()
+                        counts[l] += 1
+                imgs_batch.clear()
+                labels_batch_multi.clear()
+            if all(counts[l] >= max_per_class for l in range(num_classes)):
+                break
+        if imgs_batch:
+            batch = torch.cat(imgs_batch, dim=0).to(device)
+            feats = model.encode_image(batch)
+            feats = feats / (feats.norm(dim=-1, keepdim=True) + 1e-12)
+            if D is None:
+                D = int(feats.shape[1])
+            for f, ls in zip(feats, labels_batch_multi):
+                for l in ls:
+                    if counts[l] >= max_per_class:
+                        continue
+                    if l not in sums:
+                        sums[l] = f.detach().clone()
+                    else:
+                        sums[l] = sums[l] + f.detach()
+                    counts[l] += 1
+            imgs_batch.clear()
+            labels_batch_multi.clear()
+    class_vecs: List[torch.Tensor] = []
+    for lid in range(num_classes):
+        c = counts.get(lid, 0)
+        if c <= 0:
+            class_vecs.append(torch.zeros(int(D or 0), device=device))
+        else:
+            m = sums[lid] / float(c)
+            m = m / (m.norm() + 1e-12)
+            class_vecs.append(m)
+    if not class_vecs:
+        return torch.zeros((0, int(D or 0)), device=device)
+    return torch.stack(class_vecs, dim=0)
+
+
+def build_visual_class_prototypes_subset(model, device, ds_wrapper, label_ids: List[int], max_per_class: int, batch_size: int, preprocess_eval) -> torch.Tensor:
+    D = None
+    sums: Dict[int, torch.Tensor] = {}
+    counts: Dict[int, int] = {lid: 0 for lid in label_ids}
+    label_set = set(label_ids)
+    imgs_batch: List[torch.Tensor] = []
+    labels_batch_multi: List[List[int]] = []
+    with torch.no_grad():
+        for i in range(len(ds_wrapper)):
+            img, lab = ds_wrapper[i]
+            lab_all = [lid for lid in normalize_labels_to_list(lab) if lid in label_set and counts.get(lid, 0) < max_per_class]
+            if not lab_all:
+                continue
+            tensor_img = img if isinstance(img, torch.Tensor) else preprocess_eval(img)
+            imgs_batch.append(tensor_img.unsqueeze(0))
+            labels_batch_multi.append(lab_all)
+            if len(imgs_batch) >= max(1, batch_size):
+                batch = torch.cat(imgs_batch, dim=0).to(device)
+                feats = model.encode_image(batch)
+                feats = feats / (feats.norm(dim=-1, keepdim=True) + 1e-12)
+                if D is None:
+                    D = int(feats.shape[1])
+                for f, ls in zip(feats, labels_batch_multi):
+                    for l in ls:
+                        if counts[l] >= max_per_class:
+                            continue
+                        if l not in sums:
+                            sums[l] = f.detach().clone()
+                        else:
+                            sums[l] = sums[l] + f.detach()
+                        counts[l] += 1
+                imgs_batch.clear()
+                labels_batch_multi.clear()
+            if all(counts[lid] >= max_per_class for lid in label_ids):
+                break
+        if imgs_batch:
+            batch = torch.cat(imgs_batch, dim=0).to(device)
+            feats = model.encode_image(batch)
+            feats = feats / (feats.norm(dim=-1, keepdim=True) + 1e-12)
+            if D is None:
+                D = int(feats.shape[1])
+            for f, ls in zip(feats, labels_batch_multi):
+                for l in ls:
+                    if counts[l] >= max_per_class:
+                        continue
+                    if l not in sums:
+                        sums[l] = f.detach().clone()
+                    else:
+                        sums[l] = sums[l] + f.detach()
+                    counts[l] += 1
+            imgs_batch.clear()
+            labels_batch_multi.clear()
+    class_vecs: List[torch.Tensor] = []
+    for lid in label_ids:
+        c = counts.get(lid, 0)
+        if c <= 0:
+            class_vecs.append(torch.zeros(int(D or 0), device=device))
+        else:
+            m = sums[lid] / float(c)
+            m = m / (m.norm() + 1e-12)
+            class_vecs.append(m)
+    if not class_vecs:
+        return torch.zeros((0, int(D or 0)), device=device)
+    return torch.stack(class_vecs, dim=0)
+
+
+def cosine_distance_matrix(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+    S = (A @ B.t()).clamp(-1.0, 1.0)
+    return 1.0 - S
+
+
+def directed_covering_distance(T_src: torch.Tensor, V_src: torch.Tensor,
+                               T_tgt: torch.Tensor, V_tgt: torch.Tensor) -> float:
+    if V_src.size(0) == 0 or V_tgt.size(0) == 0:
+        return 0.0
+    F_src = torch.cat([V_src, T_src], dim=1)
+    F_tgt = torch.cat([V_tgt, T_tgt], dim=1)
+    F_src = F_src / (F_src.norm(dim=-1, keepdim=True) + 1e-12)
+    F_tgt = F_tgt / (F_tgt.norm(dim=-1, keepdim=True) + 1e-12)
+    C = cosine_distance_matrix(F_src, F_tgt)
+    d_row_min = C.min(dim=1).values
+    return float(d_row_min.mean().item())
+
+
+def tasks_to_upstream_similarity_matrix(vision_model, device,
+                                        upstream_names: List[str], name_to_idx: Dict[str, int],
+                                        dataset_list, classes_names_list, templates_list,
+                                        downstream_name: str, cil_splits: int,
+                                        preprocess_eval,
+                                        max_images_per_class: int, vision_batch_size: int,
+                                        dataset_root: str,
+                                        beta: float,
+                                        clip_text_model=None,
+                                        seed: int = 32) -> Tuple[np.ndarray, List[str]]:
+    T_up_list: List[torch.Tensor] = []
+    V_up_list: List[torch.Tensor] = []
+    for nm in upstream_names:
+        i = name_to_idx[nm]
+        classnames = classes_names_list[i]
+        templates = templates_list[i]
+        ds_wrapper = dataset_list[i]
+        K = len(classnames)
+        T_k = build_text_class_prototypes(clip_text_model, device, classnames, templates, batch_size=256)
+        V_k = build_visual_class_prototypes(vision_model, device, ds_wrapper, K, max_per_class=int(max_images_per_class), batch_size=int(vision_batch_size), preprocess_eval=preprocess_eval)
+        T_up_list.append(T_k)
+        V_up_list.append(V_k)
+
+    ds_idx = name_to_idx[downstream_name]
+    cfg_ds = type("Cfg", (), {})()
+    cfg_ds.dataset = "MTIL"
+    cfg_ds.dataset_root = dataset_root
+    cfg_ds.MTIL_order_2 = False
+    cfg_ds.train_one_dataset = ds_idx
+    cfg_ds.seed = int(seed)
+    cfg_ds.use_validation = False
+    train_list, train_classes_names, train_templates, _ = get_mtil_dataset(
+        cfg_ds, 'train', transforms=preprocess_eval
+    )
+    test_list, _, _, _ = get_mtil_dataset(
+        cfg_ds, 'test', transforms=preprocess_eval
+    )
+    assert len(train_list) == 1 and len(test_list) == 1, "Expected single selected dataset for downstream"
+    classnames_single = train_classes_names[0]
+    templates_single = train_templates[0]
+    _, _, _, class_ids_per_task, _ = build_mtil_cil_scenarios(
+        train_list[0], test_list[0], classnames_single, cil_splits, seed=int(seed)
+    )
+
+    distances = np.zeros((cil_splits, len(upstream_names)), dtype=float)
+    for t, cls_ids in enumerate(class_ids_per_task):
+        cls_names_t = [classnames_single[c] for c in cls_ids]
+        T_t = build_text_class_prototypes(clip_text_model, device, cls_names_t, templates_single, batch_size=256)
+        ds_down = train_list[0]
+        V_t = build_visual_class_prototypes_subset(vision_model, device, ds_down, list(cls_ids), max_per_class=int(max_images_per_class), batch_size=int(vision_batch_size), preprocess_eval=preprocess_eval)
+        for j, nm in enumerate(upstream_names):
+            distances[t, j] = directed_covering_distance(T_t, V_t, T_up_list[j], V_up_list[j])
+
+    similarity = np.exp(-float(beta) * distances)
+    similarity = np.clip(similarity, 0.0, 1.0)
+    return similarity, upstream_names
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Dataset similarity based on CLIP text and visual prototypes")
+    parser.add_argument("--dataset-root", type=str, default=os.environ.get("DATASET_ROOT", "data"))
+    parser.add_argument("--model-name", type=str, default="ViT-B/16")
+    parser.add_argument("--vision-batch-size", type=int, default=64)
+    parser.add_argument("--max-images-per-class", type=int, default=64, help="Maximum number of samples to use for each class when building visual prototypes. For efficiency, Chamfer distance is not computed using all samples in a dataset; instead, each class is sampled with at most this many examples.")
+    parser.add_argument("--beta", type=float, default=1.0, help="Similarity mapping exp(-beta * dist)")
+    parser.add_argument("--output", type=str, default="")
+    parser.add_argument("--downstream-dataset", type=str, default="", help="Name of downstream dataset (e.g., CIFAR100, PCam, FGVCAircraft, DescribableTextures)")
+    parser.add_argument("--cil-split", type=int, default=0, help="Number of CIL splits for downstream dataset; if >0, output [cil_split x 24] matrix")
+    args = parser.parse_args()
+
+    names = DEFAULT_NAMES_ORDER1
+
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    clip_model, _, preprocess_eval_clip = load_orig_clip(args.model_name, device=device, jit=False)
+    clip_model.eval()
+
+    vision_model = clip_model
+    preprocess_eval = preprocess_eval_clip
+    clip_text_model = clip_model
+
+    cfg = type("Cfg", (), {})()
+    cfg.dataset = "MTIL"
+    cfg.dataset_root = args.dataset_root
+    cfg.MTIL_order_2 = False
+    cfg.train_one_dataset = -1
+    cfg.seed = 32
+    cfg.use_validation = False
+
+    dataset_list, classes_names_list, templates_list, _ = get_mtil_dataset(
+        cfg, 'test', transforms=preprocess_eval
+    )
+
+    n = min(len(names), len(classes_names_list))
+    names = names[:n]
+    classes_names_list = classes_names_list[:n]
+    templates_list = templates_list[:n] if templates_list is not None else [None] * n
+
+    if args.downstream_dataset and int(args.cil_split) > 0:
+        repo_root = os.path.dirname(os.path.abspath(__file__))
+        alias_in = (args.downstream_dataset or '').strip()
+        alias_map = {
+            "Aircraft": "FGVCAircraft",
+            "FGVCAircraft": "FGVCAircraft",
+            "DTD": "DescribableTextures",
+            "DescribableTextures": "DescribableTextures",
+        }
+        ds_internal = alias_map.get(alias_in, alias_in) if alias_in else None
+        downstream_seed = resolve_downstream_seed(repo_root, ds_internal)
+        name_to_idx = {nm: idx for idx, nm in enumerate(names)}
+        if ds_internal not in name_to_idx:
+            raise ValueError(f"Downstream dataset '{args.downstream_dataset}' (mapped to '{ds_internal}') not found in MTIL names.")
+        cfg.seed = downstream_seed
+        dataset_list, classes_names_list, templates_list, _ = get_mtil_dataset(
+            cfg, 'test', transforms=preprocess_eval
+        )
+        n = min(len(names), len(classes_names_list))
+        names = names[:n]
+        classes_names_list = classes_names_list[:n]
+        templates_list = templates_list[:n] if templates_list is not None else [None] * n
+        upstream_names = [nm for nm in names if nm != ds_internal]
+        similarity_matrix, upstream_names = tasks_to_upstream_similarity_matrix(
+            vision_model, device,
+            upstream_names, name_to_idx,
+            dataset_list, classes_names_list, templates_list,
+            ds_internal, int(args.cil_split),
+            preprocess_eval,
+            int(args.max_images_per_class), int(args.vision_batch_size),
+            args.dataset_root, float(args.beta), clip_text_model=clip_text_model,
+            seed=downstream_seed,
+        )
+        print("DEFAULT_SIM_MATRIX = [")
+        fmt = "{:.3f}"
+        for i in range(similarity_matrix.shape[0]):
+            row_str = ", ".join(fmt.format(float(x)) for x in similarity_matrix[i])
+            print(f"    [{row_str}],")
+        print("]")
+        names_py = ", ".join([f'"{n}"' for n in upstream_names])
+        print(f"DEFAULT_SIM_UPSTREAM_NAMES = [{names_py}]")
+        if args.output:
+            import json
+            out = {
+                "downstream": args.downstream_dataset,
+                "cil_split": int(args.cil_split),
+                "upstream_names": upstream_names,
+                "similarity_matrix": similarity_matrix.tolist(),
+                "beta": float(args.beta),
+                "seed": int(downstream_seed),
+                "max_images_per_class": int(args.max_images_per_class),
+                "vision_encoder": "clip",
+                "text_encoder": "clip",
+            }
+            base, ext = os.path.splitext(args.output)
+            out2 = base + "_sim" + ext
+            with open(out2, 'w', encoding='utf-8') as f:
+                json.dump(out, f, ensure_ascii=False, indent=2)
+            print(f"Saved similarity JSON to {out2}")
+        return
+
+if __name__ == "__main__":
+    main()

class_orders/cifar100.yaml

@@ -0,0 +1 @@
+class_order: [87, 0, 52, 58, 44, 91, 68, 97, 51, 15, 94, 92, 10, 72, 49, 78, 61, 14, 8, 86, 84, 96, 18, 24, 32, 45, 88, 11, 4, 67, 69, 66, 77, 47, 79, 93, 29, 50, 57, 83, 17, 81, 41, 12, 37, 59, 25, 20, 80, 73, 1, 28, 6, 46, 62, 82, 53, 9, 31, 75, 38, 63, 33, 74, 27, 22, 36, 3, 16, 21, 60, 19, 70, 90, 89, 43, 5, 42, 65, 76, 40, 30, 23, 85, 2, 95, 56, 48, 71, 64, 98, 13, 99, 7, 34, 55, 54, 26, 35, 39]

class_orders/imagenet100.yaml

@@ -0,0 +1,11 @@
+class_order: [
+  68, 56, 78, 8, 23, 84, 90, 65, 74, 76, 40, 
+  89, 3, 92, 55, 9, 26, 80, 43, 38, 58, 70, 
+  77, 1, 85, 19, 17, 50, 28, 53, 13, 81, 45, 
+  82, 6, 59, 83, 16, 15, 44, 91, 41, 72, 60, 
+  79, 52, 20, 10, 31, 54, 37, 95, 14, 71, 96, 
+  98, 97, 2, 64, 66, 42, 22, 35, 86, 24, 34, 
+  87, 21, 99, 0, 88, 27, 18, 94, 11, 12, 47, 
+  25, 30, 46, 62, 69, 36, 61, 7, 63, 75, 5, 32, 
+  4, 51, 48, 73, 93, 39, 67, 29, 49, 57, 33
+]

class_orders/imagenet1000.yaml

@@ -0,0 +1,48 @@
+class_order: [
+  54, 7, 894, 512, 126, 337, 988, 11, 284, 493, 133, 783, 192, 979, 622, 215, 240, 548, 238, 419, 274, 108,
+  928, 856, 494, 836, 473, 650, 85, 262, 508, 590, 390, 174, 637, 288, 658, 219, 912, 142, 852, 160, 704, 289,
+  123, 323, 600, 542, 999, 634, 391, 761, 490, 842, 127, 850, 665, 990, 597, 722, 748, 14, 77, 437, 394, 859,
+  279, 539, 75, 466, 886, 312, 303, 62, 966, 413, 959, 782, 509, 400, 471, 632, 275, 730, 105, 523, 224, 186,
+  478, 507, 470, 906, 699, 989, 324, 812, 260, 911, 446, 44, 765, 759, 67, 36, 5, 30, 184, 797, 159, 741, 954,
+  465, 533, 585, 150, 101, 897, 363, 818, 620, 824, 154, 956, 176, 588, 986, 172, 223, 461, 94, 141, 621, 659,
+  360, 136, 578, 163, 427, 70, 226, 925, 596, 336, 412, 731, 755, 381, 810, 69, 898, 310, 120, 752, 93, 39,
+  326, 537, 905, 448, 347, 51, 615, 601, 229, 947, 348, 220, 949, 972, 73, 913, 522, 193, 753, 921, 257, 957,
+  691, 155, 820, 584, 948, 92, 582, 89, 379, 392, 64, 904, 169, 216, 694, 103, 410, 374, 515, 484, 624, 409,
+  156, 455, 846, 344, 371, 468, 844, 276, 740, 562, 503, 831, 516, 663, 630, 763, 456, 179, 996, 936, 248,
+  333, 941, 63, 738, 802, 372, 828, 74, 540, 299, 750, 335, 177, 822, 643, 593, 800, 459, 580, 933, 306, 378,
+  76, 227, 426, 403, 322, 321, 808, 393, 27, 200, 764, 651, 244, 479, 3, 415, 23, 964, 671, 195, 569, 917,
+  611, 644, 707, 355, 855, 8, 534, 657, 571, 811, 681, 543, 313, 129, 978, 592, 573, 128, 243, 520, 887, 892,
+  696, 26, 551, 168, 71, 398, 778, 529, 526, 792, 868, 266, 443, 24, 57, 15, 871, 678, 745, 845, 208, 188,
+  674, 175, 406, 421, 833, 106, 994, 815, 581, 676, 49, 619, 217, 631, 934, 932, 568, 353, 863, 827, 425, 420,
+  99, 823, 113, 974, 438, 874, 343, 118, 340, 472, 552, 937, 0, 10, 675, 316, 879, 561, 387, 726, 255, 407,
+  56, 927, 655, 809, 839, 640, 297, 34, 497, 210, 606, 971, 589, 138, 263, 587, 993, 973, 382, 572, 735, 535,
+  139, 524, 314, 463, 895, 376, 939, 157, 858, 457, 935, 183, 114, 903, 767, 666, 22, 525, 902, 233, 250, 825,
+  79, 843, 221, 214, 205, 166, 431, 860, 292, 976, 739, 899, 475, 242, 961, 531, 110, 769, 55, 701, 532, 586,
+  729, 253, 486, 787, 774, 165, 627, 32, 291, 962, 922, 222, 705, 454, 356, 445, 746, 776, 404, 950, 241, 452,
+  245, 487, 706, 2, 137, 6, 98, 647, 50, 91, 202, 556, 38, 68, 649, 258, 345, 361, 464, 514, 958, 504, 826,
+  668, 880, 28, 920, 918, 339, 315, 320, 768, 201, 733, 575, 781, 864, 617, 171, 795, 132, 145, 368, 147, 327,
+  713, 688, 848, 690, 975, 354, 853, 148, 648, 300, 436, 780, 693, 682, 246, 449, 492, 162, 97, 59, 357, 198,
+  519, 90, 236, 375, 359, 230, 476, 784, 117, 940, 396, 849, 102, 122, 282, 181, 130, 467, 88, 271, 793, 151,
+  847, 914, 42, 834, 521, 121, 29, 806, 607, 510, 837, 301, 669, 78, 256, 474, 840, 52, 505, 547, 641, 987,
+  801, 629, 491, 605, 112, 429, 401, 742, 528, 87, 442, 910, 638, 785, 264, 711, 369, 428, 805, 744, 380, 725,
+  480, 318, 997, 153, 384, 252, 985, 538, 654, 388, 100, 432, 832, 565, 908, 367, 591, 294, 272, 231, 213,
+  196, 743, 817, 433, 328, 970, 969, 4, 613, 182, 685, 724, 915, 311, 931, 865, 86, 119, 203, 268, 718, 317,
+  926, 269, 161, 209, 807, 645, 513, 261, 518, 305, 758, 872, 58, 65, 146, 395, 481, 747, 41, 283, 204, 564,
+  185, 777, 33, 500, 609, 286, 567, 80, 228, 683, 757, 942, 134, 673, 616, 960, 450, 350, 544, 830, 736, 170,
+  679, 838, 819, 485, 430, 190, 566, 511, 482, 232, 527, 411, 560, 281, 342, 614, 662, 47, 771, 861, 692, 686,
+  277, 373, 16, 946, 265, 35, 9, 884, 909, 610, 358, 18, 737, 977, 677, 803, 595, 135, 458, 12, 46, 418, 599,
+  187, 107, 992, 770, 298, 104, 351, 893, 698, 929, 502, 273, 20, 96, 791, 636, 708, 267, 867, 772, 604, 618,
+  346, 330, 554, 816, 664, 716, 189, 31, 721, 712, 397, 43, 943, 804, 296, 109, 576, 869, 955, 17, 506, 963,
+  786, 720, 628, 779, 982, 633, 891, 734, 980, 386, 365, 794, 325, 841, 878, 370, 695, 293, 951, 66, 594, 717,
+  116, 488, 796, 983, 646, 499, 53, 1, 603, 45, 424, 875, 254, 237, 199, 414, 307, 362, 557, 866, 341, 19,
+  965, 143, 555, 687, 235, 790, 125, 173, 364, 882, 727, 728, 563, 495, 21, 558, 709, 719, 877, 352, 83, 998,
+  991, 469, 967, 760, 498, 814, 612, 715, 290, 72, 131, 259, 441, 924, 773, 48, 625, 501, 440, 82, 684, 862,
+  574, 309, 408, 680, 623, 439, 180, 652, 968, 889, 334, 61, 766, 399, 598, 798, 653, 930, 149, 249, 890, 308,
+  881, 40, 835, 577, 422, 703, 813, 857, 995, 602, 583, 167, 670, 212, 751, 496, 608, 84, 639, 579, 178, 489,
+  37, 197, 789, 530, 111, 876, 570, 700, 444, 287, 366, 883, 385, 536, 460, 851, 81, 144, 60, 251, 13, 953,
+  270, 944, 319, 885, 710, 952, 517, 278, 656, 919, 377, 550, 207, 660, 984, 447, 553, 338, 234, 383, 749,
+  916, 626, 462, 788, 434, 714, 799, 821, 477, 549, 661, 206, 667, 541, 642, 689, 194, 152, 981, 938, 854,
+  483, 332, 280, 546, 389, 405, 545, 239, 896, 672, 923, 402, 423, 907, 888, 140, 870, 559, 756, 25, 211, 158,
+  723, 635, 302, 702, 453, 218, 164, 829, 247, 775, 191, 732, 115, 331, 901, 416, 873, 754, 900, 435, 762,
+  124, 304, 329, 349, 295, 95, 451, 285, 225, 945, 697, 417
+]

class_orders/tinyimagenet.yaml

@@ -0,0 +1,17 @@
+class_order: [
+  131, 181, 22, 172, 144, 92, 97, 187, 58, 93, 6, 70, 106, 68, 
+  153, 168, 179, 199, 29, 46, 9, 142, 134, 88, 193, 110, 26, 
+  32, 117, 112, 17, 39, 166, 13, 94, 138, 109, 147, 51, 101, 
+  59, 188, 116, 5, 170, 99, 100, 167, 180, 146, 65, 1, 104, 
+  43, 38, 184, 123, 171, 137, 162, 71, 44, 95, 174, 12, 7, 
+  54, 152, 21, 47, 28, 176, 34, 2, 132, 118, 42, 189, 150, 
+  14, 165, 41, 192, 45, 82, 128, 63, 57, 197, 160, 53, 75, 
+  108, 135, 121, 159, 183, 67, 169, 50, 87, 69, 89, 196, 
+  115, 19, 148, 96, 86, 11, 8, 60, 33, 173, 78, 4, 119, 105, 
+  182, 127, 177, 30, 186, 40, 49, 178, 76, 157, 161, 73, 164, 
+  151, 31, 74, 191, 27, 125, 198, 81, 20, 155, 114, 139, 36, 
+  61, 56, 145, 48, 16, 83, 62, 85, 126, 0, 102, 23, 3, 140, 
+  15, 195, 133, 113, 190, 141, 52, 163, 156, 80, 111, 90, 175, 
+  143, 120, 84, 18, 25, 79, 37, 154, 136, 64, 158, 24, 185, 
+  72, 35, 129, 55, 149, 91, 122, 77, 103, 124, 130, 66, 10, 107, 194, 98
+]

clip/README.md

@@ -0,0 +1 @@
+This folder is a lightly modified version of https://github.com/openai/CLIP.

clip/__init__.py

@@ -0,0 +1 @@
+from .clip import *

clip/adapter.py

@@ -0,0 +1,69 @@
+# --------------------------------------------------------
+# References:
+# https://github.com/jxhe/unify-parameter-efficient-tuning
+# --------------------------------------------------------
+
+import math
+import torch
+import torch.nn as nn
+
+
+class Adapter(nn.Module):
+    def __init__(self,
+                 d_model=None,
+                 bottleneck=None,
+                 dropout=0.0,
+                 init_option="lora",
+                 adapter_scalar="1.0",
+                 adapter_layernorm_option="in"):
+        super().__init__()
+        self.n_embd = d_model if d_model is None else d_model
+        self.down_size = bottleneck
+
+        #_before
+        self.adapter_layernorm_option = adapter_layernorm_option
+
+        self.adapter_layer_norm_before = None
+        if adapter_layernorm_option == "in" or adapter_layernorm_option == "out":
+            self.adapter_layer_norm_before = nn.LayerNorm(self.n_embd)
+
+        if adapter_scalar == "learnable_scalar":
+            self.scale = nn.Parameter(torch.ones(1))
+        else:
+            self.scale = float(adapter_scalar)
+
+        self.down_proj = nn.Linear(self.n_embd, 64)
+        self.non_linear_func = nn.ReLU()
+        self.up_proj = nn.Linear(self.down_size, self.n_embd)
+
+        self.dropout = dropout
+        if init_option == "bert":
+            raise NotImplementedError
+        elif init_option == "lora":
+            with torch.no_grad():
+                nn.init.kaiming_uniform_(self.down_proj.weight, a=math.sqrt(5))
+                nn.init.zeros_(self.up_proj.weight)
+                nn.init.zeros_(self.down_proj.bias)
+                nn.init.zeros_(self.up_proj.bias)
+
+    def forward(self, x, add_residual=True, residual=None):
+
+        residual = x if residual is None else residual
+        if self.adapter_layernorm_option == 'in': #  none
+            x = self.adapter_layer_norm_before(x)
+
+        down = self.down_proj(x)
+        down = self.non_linear_func(down)
+        down = nn.functional.dropout(down, p=self.dropout, training=self.training)
+        up = self.up_proj(down)
+
+        up = up * self.scale
+
+        if self.adapter_layernorm_option == 'out': #  none
+            up = self.adapter_layer_norm_before(up)
+
+        if add_residual:
+            output = up + residual
+        else:
+            output = up
+        return output

clip/bpe_simple_vocab_16e6.txt.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
+size 1356917

clip/clip.py

@@ -0,0 +1,310 @@
+# Code ported from https://github.com/openai/CLIP
+
+import hashlib
+import os
+import urllib
+import warnings
+from typing import Union, List
+
+import torch
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize, RandomResizedCrop, InterpolationMode
+from tqdm import tqdm
+
+from clip.model import build_model
+from clip.tokenizer import SimpleTokenizer as _Tokenizer
+
+__all__ = ["available_models", "load", "tokenize"]
+_tokenizer = _Tokenizer()
+
+_MODELS = {
+    "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt",
+    "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt",
+    "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt",
+    "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt",
+    "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
+    "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt",
+}
+
+
+def _download(url: str, root: str = os.path.expanduser("~/.cache/clip")):
+    os.makedirs(root, exist_ok=True)
+    filename = os.path.basename(url)
+
+    expected_sha256 = url.split("/")[-2]
+    download_target = os.path.join(root, filename)
+
+    if os.path.exists(download_target) and not os.path.isfile(download_target):
+        raise RuntimeError(f"{download_target} exists and is not a regular file")
+
+    if os.path.isfile(download_target):
+        if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256:
+            return download_target
+        else:
+            warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+    if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256:
+        raise RuntimeError(f"Model has been downloaded but the SHA256 checksum does not not match")
+
+    return download_target
+
+def _convert_to_rgb(image):
+    return image.convert('RGB')
+
+def _transform(n_px: int, is_train: bool):
+    normalize = Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
+    if is_train:
+        return Compose([
+            RandomResizedCrop(n_px, scale=(0.9, 1.0), interpolation=InterpolationMode.BICUBIC),
+            _convert_to_rgb,
+            ToTensor(),
+            normalize,
+        ])
+    else:
+        return Compose([
+            Resize(n_px, interpolation=InterpolationMode.BICUBIC),
+            CenterCrop(n_px),
+            _convert_to_rgb,
+            ToTensor(),
+            normalize,
+        ])
+
+
+
+def available_models() -> List[str]:
+    """Returns the names of available CLIP models"""
+    return list(_MODELS.keys())
+
+# def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False, download_root: str = None):
+#     """Load a CLIP model
+
+#     Parameters
+#     ----------
+#     name : str
+#         A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
+
+#     device : Union[str, torch.device]
+#         The device to put the loaded model
+
+#     jit : bool
+#         Whether to load the optimized JIT model or more hackable non-JIT model (default).
+
+#     download_root: str
+#         path to download the model files; by default, it uses "~/.cache/clip"
+
+#     Returns
+#     -------
+#     model : torch.nn.Module
+#         The CLIP model
+
+#     preprocess : Callable[[PIL.Image], torch.Tensor]
+#         A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
+#     """
+#     if name in _MODELS:
+#         model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip"))
+#     elif os.path.isfile(name):
+#         model_path = name
+#     else:
+#         raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
+
+#     try:
+#         # loading JIT archive
+#         model = torch.jit.load(model_path, map_location=device if jit else "cpu").eval()
+#         state_dict = None
+#     except RuntimeError:
+#         # loading saved state dict
+#         if jit:
+#             warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
+#             jit = False
+#         state_dict = torch.load(model_path, map_location="cpu")
+
+#     if not jit:
+#         model = build_model(state_dict or model.state_dict()).to(device)
+#         if str(device) == "cpu":
+#             model.float()
+#         return model, _transform(model.visual.input_resolution)
+
+#     # patch the device names
+#     device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
+#     device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
+
+#     def patch_device(module):
+#         try:
+#             graphs = [module.graph] if hasattr(module, "graph") else []
+#         except RuntimeError:
+#             graphs = []
+
+#         if hasattr(module, "forward1"):
+#             graphs.append(module.forward1.graph)
+
+#         for graph in graphs:
+#             for node in graph.findAllNodes("prim::Constant"):
+#                 if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"):
+#                     node.copyAttributes(device_node)
+
+#     model.apply(patch_device)
+#     patch_device(model.encode_image)
+#     patch_device(model.encode_text)
+
+#     # patch dtype to float32 on CPU
+#     if str(device) == "cpu":
+#         float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
+#         float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
+#         float_node = float_input.node()
+
+#         def patch_float(module):
+#             try:
+#                 graphs = [module.graph] if hasattr(module, "graph") else []
+#             except RuntimeError:
+#                 graphs = []
+
+#             if hasattr(module, "forward1"):
+#                 graphs.append(module.forward1.graph)
+
+#             for graph in graphs:
+#                 for node in graph.findAllNodes("aten::to"):
+#                     inputs = list(node.inputs())
+#                     for i in [1, 2]:  # dtype can be the second or third argument to aten::to()
+#                         if inputs[i].node()["value"] == 5:
+#                             inputs[i].node().copyAttributes(float_node)
+
+#         model.apply(patch_float)
+#         patch_float(model.encode_image)
+#         patch_float(model.encode_text)
+
+#         model.float()
+
+#     return model, _transform(model.input_resolution.item())    
+
+
+def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit=True, is_train=False, pretrained=True):
+    """Load a CLIP model
+    Parameters
+    ----------
+    name : str
+        A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
+    device : Union[str, torch.device]
+        The device to put the loaded model
+    jit : bool
+        Whether to load the optimized JIT model (default) or more hackable non-JIT model.
+    Returns
+    -------
+    model : torch.nn.Module
+        The CLIP model
+    preprocess : Callable[[PIL.Image], torch.Tensor]
+        A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
+    """
+    if name in _MODELS:
+        model_path = _download(_MODELS[name])
+    elif os.path.isfile(name):
+        model_path = name
+    else:
+        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
+
+    try:
+        # loading JIT archive
+        model = torch.jit.load(model_path, map_location=device if jit else "cpu").eval()
+        state_dict = None
+    except RuntimeError:
+        # loading saved state dict
+        if jit:
+            warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
+            jit = False
+        state_dict = torch.load(model_path, map_location="cpu")
+
+    if not jit:
+        try:
+            model = build_model(state_dict or model.state_dict()).to(device)
+        except KeyError:
+            sd = {k[7:]: v for k,v in state_dict["state_dict"].items()}
+            model = build_model(sd).to(device)
+
+        if str(device) == "cpu":
+            model.float()
+        return model, \
+               _transform(model.visual.input_resolution, is_train=True), \
+               _transform(model.visual.input_resolution, is_train=False)
+
+    # patch the device names
+    device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
+    device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
+
+    def patch_device(module):
+        graphs = [module.graph] if hasattr(module, "graph") else []
+        if hasattr(module, "forward1"):
+            graphs.append(module.forward1.graph)
+
+        for graph in graphs:
+            for node in graph.findAllNodes("prim::Constant"):
+                if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"):
+                    node.copyAttributes(device_node)
+
+    model.apply(patch_device)
+    patch_device(model.encode_image)
+    patch_device(model.encode_text)
+
+    # patch dtype to float32 on CPU
+    if str(device) == "cpu":
+        float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
+        float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
+        float_node = float_input.node()
+
+        def patch_float(module):
+            graphs = [module.graph] if hasattr(module, "graph") else []
+            if hasattr(module, "forward1"):
+                graphs.append(module.forward1.graph)
+
+            for graph in graphs:
+                for node in graph.findAllNodes("aten::to"):
+                    inputs = list(node.inputs())
+                    for i in [1, 2]:  # dtype can be the second or third argument to aten::to()
+                        if inputs[i].node()["value"] == 5:
+                            inputs[i].node().copyAttributes(float_node)
+
+        model.apply(patch_float)
+        patch_float(model.encode_image)
+        patch_float(model.encode_text)
+
+        model.float()
+
+    return model, \
+           _transform(model.input_resolution.item(), is_train=True), \
+           _transform(model.input_resolution.item(), is_train=False)
+
+
+def tokenize(texts: Union[str, List[str]], context_length: int = 77) -> torch.LongTensor:
+    """
+    Returns the tokenized representation of given input string(s)
+    Parameters
+    ----------
+    texts : Union[str, List[str]]
+        An input string or a list of input strings to tokenize
+    context_length : int
+        The context length to use; all CLIP models use 77 as the context length
+    Returns
+    -------
+    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]
+    """
+    if isinstance(texts, str):
+        texts = [texts]
+
+    sot_token = _tokenizer.encoder["<start_of_text>"]
+    eot_token = _tokenizer.encoder["<end_of_text>"]
+    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
+    result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+    for i, tokens in enumerate(all_tokens):
+        if len(tokens) > context_length: # Truncate
+            tokens = tokens[:context_length]
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result

clip/model.py

@@ -0,0 +1,713 @@
+from collections import OrderedDict
+from typing import Tuple, Union
+
+import os
+import json
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from .adapter import Adapter
+from torch.distributions.normal import Normal
+from collections import Counter
+
+global_taskid = 0
+global_is_train=True
+class SparseDispatcher(object):
+    """Helper for implementing a mixture of experts.
+    The purpose of this class is to create input minibatches for the
+    experts and to combine the results of the experts to form a unified
+    output tensor.
+    There are two functions:
+    dispatch - take an input Tensor and create input Tensors for each expert.
+    combine - take output Tensors from each expert and form a combined output
+      Tensor.  Outputs from different experts for the same batch element are
+      summed together, weighted by the provided "gates".
+    The class is initialized with a "gates" Tensor, which specifies which
+    batch elements go to which experts, and the weights to use when combining
+    the outputs.  Batch element b is sent to expert e iff gates[b, e] != 0.
+    The inputs and outputs are all two-dimensional [batch, depth].
+    Caller is responsible for collapsing additional dimensions prior to
+    calling this class and reshaping the output to the original shape.
+    See common_layers.reshape_like().
+    Example use:
+    gates: a float32 `Tensor` with shape `[batch_size, num_experts]`
+    inputs: a float32 `Tensor` with shape `[batch_size, input_size]`
+    experts: a list of length `num_experts` containing sub-networks.
+    dispatcher = SparseDispatcher(num_experts, gates)
+    expert_inputs = dispatcher.dispatch(inputs)
+    expert_outputs = [experts[i](expert_inputs[i]) for i in range(num_experts)]
+    outputs = dispatcher.combine(expert_outputs)
+    The preceding code sets the output for a particular example b to:
+    output[b] = Sum_i(gates[b, i] * experts[i](inputs[b]))
+    This class takes advantage of sparsity in the gate matrix by including in the
+    `Tensor`s for expert i only the batch elements for which `gates[b, i] > 0`.
+    """
+
+    def __init__(self, num_experts, gates):
+        """Create a SparseDispatcher."""
+
+        self._gates = gates
+        self._num_experts = num_experts
+
+        sorted_experts, index_sorted_experts = torch.nonzero(gates).sort(0)
+
+        # drop indices
+        _, self._expert_index = sorted_experts.split(1, dim=1)
+        # get according batch index for each expert
+        self._batch_index = torch.nonzero(gates)[index_sorted_experts[:, 1], 0]
+        # calculate num samples that each expert gets
+        self._part_sizes = (gates > 0).sum(0).tolist()
+        # expand gates to match with self._batch_index
+        gates_exp = gates[self._batch_index.flatten()]
+        self._nonzero_gates = torch.gather(gates_exp, 1, self._expert_index)
+
+    def dispatch(self, inp):
+        """Create one input Tensor for each expert.
+        The `Tensor` for a expert `i` contains the slices of `inp` corresponding
+        to the batch elements `b` where `gates[b, i] > 0`.
+        Args:
+          inp: a `Tensor` of shape "[batch_size, <extra_input_dims>]`
+        Returns:
+          a list of `num_experts` `Tensor`s with shapes
+            `[expert_batch_size_i, <extra_input_dims>]`.
+        """
+
+        # assigns samples to experts whose gate is nonzero
+
+        inp_exp = inp[self._batch_index].squeeze(1)
+        return torch.split(inp_exp, self._part_sizes, dim=0)
+
+    def combine(self, expert_out, multiply_by_gates=True):
+        """Sum together the expert output, weighted by the gates.
+        The slice corresponding to a particular batch element `b` is computed
+        as the sum over all experts `i` of the expert output, weighted by the
+        corresponding gate values.  If `multiply_by_gates` is set to False, the
+        gate values are ignored.
+        Args:
+          expert_out: a list of `num_experts` `Tensor`s, each with shape
+            `[expert_batch_size_i, <extra_output_dims>]`.
+          multiply_by_gates: a boolean
+        Returns:
+          a `Tensor` with shape `[batch_size, <extra_output_dims>]`.
+        """
+        # apply exp to expert outputs, so we are not longer in log space
+
+        stitched = torch.cat(expert_out, 0)
+        if multiply_by_gates:
+            stitched = stitched.mul(self._nonzero_gates)  # 加权
+
+        zeros = torch.zeros(self._gates.size(0), expert_out[-1].size(1), device=stitched.device)
+        # combine samples that have been processed by the same k experts
+
+        combined = zeros.index_add(0, self._batch_index, stitched.float())
+        # add eps to all zero values in order to avoid nans when going back to log space
+        # back to log space
+        return combined
+
+    def expert_to_gates(self):
+        """Gate values corresponding to the examples in the per-expert `Tensor`s.
+        Returns:
+          a list of `num_experts` one-dimensional `Tensor`s with type `tf.float32`
+              and shapes `[expert_batch_size_i]`
+        """
+        # split nonzero gates for each expert
+        return torch.split(self._nonzero_gates, self._part_sizes, dim=0)
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.relu(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x, key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+
+        return x[0]
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_resolution = input_resolution
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.avgpool = nn.AvgPool2d(2)
+        self.relu = nn.ReLU(inplace=True)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        def stem(x):
+            for conv, bn in [(self.conv1, self.bn1), (self.conv2, self.bn2), (self.conv3, self.bn3)]:
+                x = self.relu(bn(conv(x)))
+            x = self.avgpool(x)
+            return x
+
+        x = x.type(self.conv1.weight.dtype)
+        x = stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.attnpool(x)
+
+        return x
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None, text_or_image=None):
+        super().__init__()
+        self.register_buffer("mean", torch.tensor([0.0]))
+        self.register_buffer("std", torch.tensor([1.0]))
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+        self.is_train = global_is_train
+        self.step = 1
+        self.top_k = 2
+        self.ffn_num = 64
+        self.experts_num = 2
+        self.softmax = nn.Softmax(1)
+        self.softplus = nn.Softplus()
+        self.noisy_gating = True
+        self.adaptmlp_list = nn.ModuleList()
+        self.text_or_image = text_or_image
+        if text_or_image == 'text':
+            # print('text transformer')
+            self.choose_map_text = torch.zeros([ self.experts_num])
+        else:
+            # print('image transformer')
+            self.choose_map_image = torch.zeros([ self.experts_num])
+        self.router_list = nn.ParameterList()
+        self.w_noise_list = nn.ParameterList()
+        for i in range(self.step):
+            self.router_list.append(nn.Parameter(torch.zeros(d_model, self.experts_num), requires_grad=True))
+            self.w_noise_list.append(nn.Parameter(torch.zeros(d_model, self.experts_num), requires_grad=True))
+        for i in range(self.experts_num):  #
+            self.adaptmlp = Adapter(d_model=d_model, dropout=0.1, bottleneck=self.ffn_num,
+                                    init_option='lora',
+                                    adapter_scalar=0.1,
+                                    adapter_layernorm_option='none',
+                                    )
+            self.adaptmlp_list.append(self.adaptmlp)
+
+        # self.taskid = None
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def cv_squared(self, x):
+        """The squared coefficient of variation of a sample.
+        Useful as a loss to encourage a positive distribution to be more uniform.
+        Epsilons added for numerical stability.
+        Returns 0 for an empty Tensor.
+        Args:
+        x: a `Tensor`.
+        Returns:
+        a `Scalar`.
+        """
+        eps = 1e-10
+        # if only num_experts = 1
+
+        if x.shape[0] == 1:
+            return torch.tensor([0], device=x.device, dtype=x.dtype)
+        return x.float().var() / (x.float().mean()**2 + eps)
+
+    def _gates_to_load(self, gates):
+        """Compute the true load per expert, given the gates.
+        The load is the number of examples for which the corresponding gate is >0.
+        Args:
+        gates: a `Tensor` of shape [batch_size, n]
+        Returns:
+        a float32 `Tensor` of shape [n]
+        """
+        return (gates > 0).sum(0)
+
+    def _prob_in_top_k(self, clean_values, noisy_values, noise_stddev, noisy_top_values):
+        """Helper function to NoisyTopKGating.
+        Computes the probability that value is in top k, given different random noise.
+        This gives us a way of backpropagating from a loss that balances the number
+        of times each expert is in the top k experts per example.
+        In the case of no noise, pass in None for noise_stddev, and the result will
+        not be differentiable.
+        Args:
+        clean_values: a `Tensor` of shape [batch, n].
+        noisy_values: a `Tensor` of shape [batch, n].  Equal to clean values plus
+          normally distributed noise with standard deviation noise_stddev.
+        noise_stddev: a `Tensor` of shape [batch, n], or None
+        noisy_top_values: a `Tensor` of shape [batch, m].
+           "values" Output of tf.top_k(noisy_top_values, m).  m >= k+1
+        Returns:
+        a `Tensor` of shape [batch, n].
+        """
+        # print('1231',clean_values)  # 全nan
+        batch = clean_values.size(0)
+        m = noisy_top_values.size(1)
+        top_values_flat = noisy_top_values.flatten()
+
+        threshold_positions_if_in = torch.arange(batch, device=clean_values.device) * m + self.top_k
+        threshold_if_in = torch.unsqueeze(torch.gather(top_values_flat, 0, threshold_positions_if_in), 1)
+        is_in = torch.gt(noisy_values, threshold_if_in)
+        threshold_positions_if_out = threshold_positions_if_in - 1
+        threshold_if_out = torch.unsqueeze(torch.gather(top_values_flat, 0, threshold_positions_if_out), 1)
+        # is each value currently in the top k.
+        normal = Normal(self.mean, self.std)
+        #
+
+        prob_if_in = normal.cdf((clean_values - threshold_if_in)/noise_stddev)
+        prob_if_out = normal.cdf((clean_values - threshold_if_out)/noise_stddev)
+        prob = torch.where(is_in, prob_if_in, prob_if_out)
+        return prob
+
+    def noisy_top_k_gating(self, x, train, w_gate, w_noise, noise_epsilon=1e-2):
+        """Noisy top-k gating.
+          See paper: https://arxiv.org/abs/1701.06538.
+          Args:
+            x: input Tensor with shape [batch_size, input_size]
+            train: a boolean - we only add noise at training time.
+            noise_epsilon: a float
+          Returns:
+            gates: a Tensor with shape [batch_size, num_experts]
+            load: a Tensor with shape [num_experts]
+        """
+
+        clean_logits = x @ w_gate.to(x)
+        if self.noisy_gating and train:
+            raw_noise_stddev = x @ w_noise.to(x)
+            noise_stddev = ((self.softplus(raw_noise_stddev) + noise_epsilon))
+            noisy_logits = clean_logits + (torch.randn_like(clean_logits) * noise_stddev)
+            logits = noisy_logits
+        else:
+            logits = clean_logits
+        # calculate topk + 1 that will be needed for the noisy gates
+        top_logits, top_indices = logits.topk(min(self.top_k + 1, self.experts_num), dim=1)
+        top_k_logits = top_logits[:, :self.top_k]
+        top_k_indices = top_indices[:, :self.top_k]
+        top_k_gates = self.softmax(top_k_logits)
+        zeros = torch.zeros_like(logits)
+        gates = zeros.scatter(1, top_k_indices, top_k_gates)
+        if self.noisy_gating and self.top_k < self.experts_num and train:  # 目前未用上
+            load = (self._prob_in_top_k(clean_logits, noisy_logits, noise_stddev, top_logits)).sum(0)
+        else:
+            load = self._gates_to_load(gates)
+        return gates, load
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.attention(self.ln_1(x))
+        if global_taskid is not None:
+            x_re = x.permute(1, 0, 2)[:, 0, :]
+            gates, load = self.noisy_top_k_gating(x_re, self.is_train, self.router_list[global_taskid],
+                                                  self.w_noise_list[global_taskid])
+            importance = gates.sum(0)
+
+            nonzero_indices = torch.nonzero(gates)
+            counter = Counter(nonzero_indices[:, 1].tolist())
+            for number, count in counter.items():
+                if self.text_or_image == 'text':
+                    self.choose_map_text[number] = self.choose_map_text[number] + count
+                else:
+                    self.choose_map_image[number] = self.choose_map_image[number] + count
+            dispatcher = SparseDispatcher(self.experts_num, gates)
+            expert_inputs = dispatcher.dispatch(x.permute(1, 0, 2).view(x.shape[1], -1))
+            expert_outputs = [self.adaptmlp_list[i](expert_inputs[i].view(expert_inputs[i].shape[0],
+                                                                          x.shape[0], x.shape[2]).to(x), add_residual=False)
+                              for i in range(self.experts_num)]
+
+            i = 0
+            while i < len(expert_outputs):
+                if expert_outputs[i].shape[0] == 0:
+                    expert_outputs.pop(i)
+                else:
+                    expert_outputs[i] = expert_outputs[i].view(expert_outputs[i].shape[0], -1)
+                    i += 1
+
+            y = dispatcher.combine(expert_outputs)
+            y = y.view(x.shape[1], x.shape[0], x.shape[2])
+            x = x + self.mlp(self.ln_2(x)) + y.permute(1, 0, 2)
+        else:
+            x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None, text_or_image=None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask, text_or_image) for _ in range(layers)])
+
+    def forward(self, x: torch.Tensor):
+        return self.resblocks(x)
+
+
+class VisualTransformer(nn.Module):
+    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int, text_or_image=None):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.output_dim = output_dim
+        # Added so this info is available. should not change anything.
+        self.patch_size = patch_size
+        self.width = width
+        self.layers = layers
+        self.heads = heads
+
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads, text_or_image=text_or_image)
+
+        self.ln_post = LayerNorm(width)
+        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+    def forward(self, x: torch.Tensor):
+        x = self.conv1(x)
+        x = x.reshape(x.shape[0], x.shape[1], -1)
+        x = x.permute(0, 2, 1)
+        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        x = self.ln_post(x[:, 0, :])
+
+        if self.proj is not None:
+            x = x @ self.proj
+
+        return x
+
+
+class CLIP(nn.Module):
+    def __init__(self,
+                 embed_dim: int,
+                 # vision
+                 image_resolution: int,
+                 vision_layers: Union[Tuple[int, int, int, int], int],
+                 vision_width: int,
+                 vision_patch_size: int,
+                 # text
+                 context_length: int,
+                 vocab_size: int,
+                 transformer_width: int,
+                 transformer_heads: int,
+                 transformer_layers: int,
+                 baseline = False
+                 ):
+        super().__init__()
+        self.baseline = baseline
+
+        self.context_length = context_length
+
+        if isinstance(vision_layers, (tuple, list)):
+            vision_heads = vision_width * 32 // 64
+            self.visual = ModifiedResNet(
+                layers=vision_layers,
+                output_dim=embed_dim,
+                heads=vision_heads,
+                input_resolution=image_resolution,
+                width=vision_width
+            )
+        else:
+            vision_heads = vision_width // 64
+            self.visual = VisualTransformer(
+                input_resolution=image_resolution,
+                patch_size=vision_patch_size,
+                width=vision_width,
+                layers=vision_layers,
+                heads=vision_heads,
+                output_dim=embed_dim,
+                text_or_image='image'
+            )
+
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask(),
+            text_or_image='text'
+        )
+
+        self.vocab_size = vocab_size
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        if isinstance(self.visual, ModifiedResNet):
+            if self.visual.attnpool is not None:
+                std = self.visual.attnpool.c_proj.in_features ** -0.5
+                nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
+
+            for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
+                for name, param in resnet_block.named_parameters():
+                    if name.endswith("bn3.weight"):
+                        nn.init.zeros_(param)
+
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    @property
+    def dtype(self):
+        return self.visual.conv1.weight.dtype
+
+    def encode_image(self, image):
+        return self.visual(image.type(self.dtype))
+
+    def encode_text(self, text):
+
+        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.type(self.dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+
+        return x
+
+    def forward(self, image, text, taskid, is_train):
+        global global_taskid, global_is_train
+        global_taskid = taskid
+        global_is_train = is_train
+        if image is None:
+            return self.encode_text(text)
+        elif text is None:
+            return self.encode_image(image)
+        image_features = self.encode_image(image)
+        text_features = self.encode_text(text)
+
+        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
+
+        # if self.baseline:
+        logit_scale = self.logit_scale.exp()
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logits_per_image.t()
+        return logits_per_image, logits_per_text
+        
+
+
+def convert_weights(model: nn.Module):
+    """Convert applicable model parameters to fp16"""
+
+    def _convert_weights_to_fp16(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.half()
+            if l.bias is not None:
+                l.bias.data = l.bias.data.half()
+
+        if isinstance(l, nn.MultiheadAttention):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.half()
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name):
+                attr = getattr(l, name)
+                if attr is not None:
+                    attr.data = attr.data.half()
+
+    model.apply(_convert_weights_to_fp16)
+
+
+def build_model(state_dict: dict):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_resolution = vision_patch_size * grid_size
+    else:
+        counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_resolution = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+
+    model = CLIP(
+        embed_dim,
+        image_resolution, vision_layers, vision_width, vision_patch_size,
+        context_length, vocab_size, transformer_width, transformer_heads, transformer_layers
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        if key in state_dict:
+            del state_dict[key]
+
+    model.load_state_dict(state_dict, strict=False)
+    for p in model.parameters():
+        p.data = p.data.float()
+    return model.eval()

clip/tokenizer.py

@@ -0,0 +1,140 @@
+import gzip
+import html
+import os
+from functools import lru_cache
+
+import ftfy
+import regex as re
+
+
+@lru_cache()
+def default_bpe():
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
+
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a corresponding list of unicode strings.
+    The reversible bpe codes work on unicode strings.
+    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+    This is a signficant percentage of your normal, say, 32K bpe vocab.
+    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+    And avoids mapping to whitespace/control characters the bpe code barfs on.
+    """
+    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8+n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """Return set of symbol pairs in a word.
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+class SimpleTokenizer(object):
+    def __init__(self, bpe_path: str = default_bpe(), special_tokens=None):
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
+        merges = merges[1:49152-256-2+1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(bytes_to_unicode().values())
+        vocab = vocab + [v+'</w>' for v in vocab]
+        for merge in merges:
+            vocab.append(''.join(merge))
+        if not special_tokens:
+            special_tokens = ['<start_of_text>', '<end_of_text>']
+        else:
+            special_tokens = ['<start_of_text>', '<end_of_text>'] + special_tokens
+        vocab.extend(special_tokens)
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {t:t for t in special_tokens}
+        special = "|".join(special_tokens)
+        self.pat = re.compile(special + r"""|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
+
+        self.vocab_size = len(self.encoder)
+        self.all_special_ids = [self.encoder[t] for t in special_tokens]
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + ( token[-1] + '</w>',)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token+'</w>'
+
+        while True:
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = ' '.join(word)
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        bpe_tokens = []
+        text = whitespace_clean(basic_clean(text)).lower()
+        for token in re.findall(self.pat, text):
+            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = ''.join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
+        return text

configs/class/aircraft.yaml

@@ -0,0 +1,76 @@
+hydra:
+  run:
+    dir: ./experiments/Aircraft
+  job:
+    chdir: true
+
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '%(message)s'
+
+class_order: ""
+dataset_root: ""
+workdir: ""
+log_path: "metrics.json"
+model_name: "ViT-B/16"
+prompt_template: "a bad photo of a {}."
+zero_shot_eval: true
+pre_task_zero_shot_eval: False
+# zs_mtil_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]  # indices of MTIL datasets (0..24) to evaluate in zero-shot; excludes any selected as downstream train_dataset
+zs_mtil_indices: [1]  # indices of MTIL datasets (0..24) to evaluate in zero-shot; excludes any selected as downstream train_dataset
+
+
+batch_size: 128
+increment: ${initial_increment}
+initial_increment: 10  
+scenario: "class"
+dataset: "cifar100"
+
+# method: "lwf"
+# lr: 7.1e-6
+# ls: 0.2
+# we: 
+# avg_freq:  
+# ref_dataset:
+# ref_sentences: 
+
+weight_decay: 0.0
+l2: 0
+ce_method: 0
+
+method: "MoE-Adapters"
+lr: 1e-3
+ls: 0.0
+#we:
+#avg_freq:
+#ref_dataset:
+#ref_sentences: random
+
+# =====================
+# DFA (two-expert + router) hyperparameters
+# =====================
+
+
+# Training schedule
+epochs_a: 1             
+epochs_b: 1             
+e2_top_k: 2             # top-k for E2 router  
+
+# Stage A (InfoNCE)
+lr_e1: 1.0e-5         
+tau_con: 0.15           # contrastive temperature for InfoNCE
+use_moco_queue: true    # enable momentum queue for more negatives
+moco_queue_size: 128  # queue size (default 4096, increases negative samples significantly)
+
+# Stage B (Task expert + Router)
+lr_e2: 5.0e-4           
+
+text_lr_e2: 1.0e-5
+lr_e2_router: 1.0e-3
+lr_top_router: 1.0e-5
+lambda_b_con: 0.001
+tau_b_con: 0.19
+num_task_experts: 2
+seed: 43

configs/class/caltech.yaml

@@ -0,0 +1,79 @@
+hydra:
+  run:
+    dir: ./experiments/Caltech101
+  job:
+    chdir: true
+
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '%(message)s'
+
+class_order: ""
+dataset_root: ""
+workdir: ""
+log_path: "metrics.json"
+model_name: "ViT-B/16"
+prompt_template: "a bad photo of a {}."
+zero_shot_eval: true
+pre_task_zero_shot_eval: False
+zs_mtil_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]  # indices of MTIL datasets (0..24) to evaluate in zero-shot; excludes any selected as downstream train_dataset
+ 
+
+
+batch_size: 128
+increment: ${initial_increment}
+initial_increment: 10  
+scenario: "class"
+dataset: "cifar100"
+
+# method: "lwf"
+# lr: 7.1e-6
+# ls: 0.2
+# we: 
+# avg_freq:  
+# ref_dataset:
+# ref_sentences: 
+
+weight_decay: 0.0
+l2: 0
+ce_method: 0
+
+method: "MoE-Adapters"
+lr: 1e-3
+ls: 0.0
+#we:
+#avg_freq:
+#ref_dataset:
+#ref_sentences: random
+
+# =====================
+# DFA (two-expert + router) hyperparameters
+# =====================
+
+
+# Training schedule
+epochs_a: 1             
+epochs_b: 1             
+e2_top_k: 2             # top-k for E2 router  
+
+# Stage A (InfoNCE)
+lr_e1: 1.0e-5         
+tau_con: 0.15           # contrastive temperature for InfoNCE
+use_moco_queue: true    # enable momentum queue for more negatives
+moco_queue_size: 128   # queue size (default 4096, increases negative samples significantly)
+ 
+
+ 
+
+# Stage B (Task expert + Router)
+lr_e2: 5.0e-4           
+
+text_lr_e2: 1.0e-5
+lr_e2_router: 1.0e-3
+lr_top_router: 1.0e-5
+lambda_b_con: 0.001
+tau_b_con: 0.19
+num_task_experts: 2
+seed: 43

configs/class/car.yaml

@@ -0,0 +1,79 @@
+hydra:
+  run:
+    dir: ./experiments/StanfordCars
+  job:
+    chdir: true
+
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '%(message)s'
+
+class_order: ""
+dataset_root: ""
+workdir: ""
+log_path: "metrics.json"
+model_name: "ViT-B/16"
+prompt_template: "a bad photo of a {}."
+zero_shot_eval: true
+pre_task_zero_shot_eval: False
+zs_mtil_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]  # indices of MTIL datasets (0..24) to evaluate in zero-shot; excludes any selected as downstream train_dataset
+ 
+
+
+batch_size: 128
+increment: ${initial_increment}
+initial_increment: 10  
+scenario: "class"
+dataset: "cifar100"
+
+# method: "lwf"
+# lr: 7.1e-6
+# ls: 0.2
+# we: 
+# avg_freq:  
+# ref_dataset:
+# ref_sentences: 
+
+weight_decay: 0.0
+l2: 0
+ce_method: 0
+
+method: "MoE-Adapters"
+lr: 1e-3
+ls: 0.0
+#we:
+#avg_freq:
+#ref_dataset:
+#ref_sentences: random
+
+# =====================
+# DFA (two-expert + router) hyperparameters
+# =====================
+
+
+# Training schedule
+epochs_a: 1             
+epochs_b: 1             
+e2_top_k: 2             # top-k for E2 router  
+
+# Stage A (InfoNCE)
+lr_e1: 1.0e-5         
+tau_con: 0.15           # contrastive temperature for InfoNCE
+use_moco_queue: true    # enable momentum queue for more negatives
+moco_queue_size: 128   # queue size (default 4096, increases negative samples significantly)
+ 
+
+ 
+
+# Stage B (Task expert + Router)
+lr_e2: 5.0e-4           
+
+text_lr_e2: 1.0e-5
+lr_e2_router: 1.0e-3
+lr_top_router: 1.0e-5
+lambda_b_con: 0.001
+tau_b_con: 0.17
+num_task_experts: 2
+seed: 43

configs/class/cifar10.yaml

@@ -0,0 +1,79 @@
+hydra:
+  run:
+    dir: ./experiments/Cifar10
+  job:
+    chdir: true
+
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '%(message)s'
+
+class_order: ""
+dataset_root: ""
+workdir: ""
+log_path: "metrics.json"
+model_name: "ViT-B/16"
+prompt_template: "a bad photo of a {}."
+zero_shot_eval: true
+pre_task_zero_shot_eval: False
+zs_mtil_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]  # indices of MTIL datasets (0..24) to evaluate in zero-shot; excludes any selected as downstream train_dataset
+ 
+
+
+batch_size: 128
+increment: ${initial_increment}
+initial_increment: 10  
+scenario: "class"
+dataset: "cifar100"
+
+# method: "lwf"
+# lr: 7.1e-6
+# ls: 0.2
+# we: 
+# avg_freq:  
+# ref_dataset:
+# ref_sentences: 
+
+weight_decay: 0.0
+l2: 0
+ce_method: 0
+
+method: "MoE-Adapters"
+lr: 1e-3
+ls: 0.0
+#we:
+#avg_freq:
+#ref_dataset:
+#ref_sentences: random
+
+# =====================
+# DFA (two-expert + router) hyperparameters
+# =====================
+
+
+# Training schedule
+epochs_a: 1             
+epochs_b: 1             
+e2_top_k: 2             # top-k for E2 router  
+
+# Stage A (InfoNCE)
+lr_e1: 1.0e-5         
+tau_con: 0.15           # contrastive temperature for InfoNCE
+use_moco_queue: true    # enable momentum queue for more negatives
+moco_queue_size: 128   # queue size (default 4096, increases negative samples significantly)
+ 
+
+ 
+
+# Stage B (Task expert + Router)
+lr_e2: 5.0e-4           
+
+text_lr_e2: 1.0e-5
+lr_e2_router: 1.0e-3
+lr_top_router: 1.0e-5
+lambda_b_con: 0.001
+tau_b_con: 0.19
+num_task_experts: 2
+seed: 43

configs/class/cifar100.yaml

@@ -0,0 +1,79 @@
+hydra:
+  run:
+    dir: ./experiments/CIFAR100
+  job:
+    chdir: true
+
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '%(message)s'
+
+class_order: ""
+dataset_root: ""
+workdir: ""
+log_path: "metrics.json"
+model_name: "ViT-B/16"
+prompt_template: "a bad photo of a {}."
+zero_shot_eval: true
+pre_task_zero_shot_eval: False
+zs_mtil_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]  # indices of MTIL datasets (0..24) to evaluate in zero-shot; excludes any selected as downstream train_dataset
+ 
+
+
+batch_size: 64
+increment: ${initial_increment}
+initial_increment: 10  
+scenario: "class"
+dataset: "cifar100"
+
+# method: "lwf"
+# lr: 7.1e-6
+# ls: 0.2
+# we: 
+# avg_freq:  
+# ref_dataset:
+# ref_sentences: 
+
+weight_decay: 0.0
+l2: 0
+ce_method: 0
+
+method: "MoE-Adapters"
+lr: 1e-3
+ls: 0.0
+#we:
+#avg_freq:
+#ref_dataset:
+#ref_sentences: random
+
+# =====================
+# DFA (two-expert + router) hyperparameters
+# =====================
+
+
+# Training schedule
+epochs_a: 1             
+epochs_b: 1             
+e2_top_k: 2             # top-k for E2 router  
+
+# Stage A (InfoNCE)
+lr_e1: 1.0e-5         
+tau_con: 0.15           # contrastive temperature for InfoNCE
+use_moco_queue: true    # enable momentum queue for more negatives
+moco_queue_size: 128   # queue size (default 4096, increases negative samples significantly)
+ 
+
+ 
+
+# Stage B (Task expert + Router)
+lr_e2: 5.0e-4           
+
+text_lr_e2: 1.0e-5
+lr_e2_router: 1.0e-3
+lr_top_router: 1.0e-5
+lambda_b_con: 0.001
+tau_b_con: 0.17
+num_task_experts: 2
+seed: 43

configs/class/clevr.yaml

@@ -0,0 +1,79 @@
+hydra:
+  run:
+    dir: ./experiments/CLEVRCount
+  job:
+    chdir: true
+
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '%(message)s'
+
+class_order: ""
+dataset_root: ""
+workdir: ""
+log_path: "metrics.json"
+model_name: "ViT-B/16"
+prompt_template: "a bad photo of a {}."
+zero_shot_eval: true
+pre_task_zero_shot_eval: False
+zs_mtil_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]  # indices of MTIL datasets (0..24) to evaluate in zero-shot; excludes any selected as downstream train_dataset
+ 
+
+
+batch_size: 128
+increment: ${initial_increment}
+initial_increment: 10  
+scenario: "class"
+dataset: "cifar100"
+
+# method: "lwf"
+# lr: 7.1e-6
+# ls: 0.2
+# we: 
+# avg_freq:  
+# ref_dataset:
+# ref_sentences: 
+
+weight_decay: 0.0
+l2: 0
+ce_method: 0
+
+method: "MoE-Adapters"
+lr: 1e-3
+ls: 0.0
+#we:
+#avg_freq:
+#ref_dataset:
+#ref_sentences: random
+
+# =====================
+# DFA (two-expert + router) hyperparameters
+# =====================
+
+
+# Training schedule
+epochs_a: 1             
+epochs_b: 1             
+e2_top_k: 2             # top-k for E2 router  
+
+# Stage A (InfoNCE)
+lr_e1: 1.0e-5         
+tau_con: 0.15           # contrastive temperature for InfoNCE
+use_moco_queue: true    # enable momentum queue for more negatives
+moco_queue_size: 128   # queue size (default 4096, increases negative samples significantly)
+ 
+
+ 
+
+# Stage B (Task expert + Router)
+lr_e2: 5.0e-4           
+
+text_lr_e2: 1.0e-5
+lr_e2_router: 1.0e-3
+lr_top_router: 1.0e-5
+lambda_b_con: 0.001
+tau_b_con: 0.19
+num_task_experts: 2
+seed: 43

configs/class/conuntry.yaml

@@ -0,0 +1,79 @@
+hydra:
+  run:
+    dir: ./experiments/Country211
+  job:
+    chdir: true
+
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '%(message)s'
+
+class_order: ""
+dataset_root: ""
+workdir: ""
+log_path: "metrics.json"
+model_name: "ViT-B/16"
+prompt_template: "a bad photo of a {}."
+zero_shot_eval: true
+pre_task_zero_shot_eval: False
+zs_mtil_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]  # indices of MTIL datasets (0..24) to evaluate in zero-shot; excludes any selected as downstream train_dataset
+ 
+
+
+batch_size: 128
+increment: ${initial_increment}
+initial_increment: 10  
+scenario: "class"
+dataset: "cifar100"
+
+# method: "lwf"
+# lr: 7.1e-6
+# ls: 0.2
+# we: 
+# avg_freq:  
+# ref_dataset:
+# ref_sentences: 
+
+weight_decay: 0.0
+l2: 0
+ce_method: 0
+
+method: "MoE-Adapters"
+lr: 1e-3
+ls: 0.0
+#we:
+#avg_freq:
+#ref_dataset:
+#ref_sentences: random
+
+# =====================
+# DFA (two-expert + router) hyperparameters
+# =====================
+
+
+# Training schedule
+epochs_a: 1             
+epochs_b: 1             
+e2_top_k: 2             # top-k for E2 router  
+
+# Stage A (InfoNCE)
+lr_e1: 1.0e-5         
+tau_con: 0.15           # contrastive temperature for InfoNCE
+use_moco_queue: true    # enable momentum queue for more negatives
+moco_queue_size: 128   # queue size (default 4096, increases negative samples significantly)
+ 
+
+ 
+
+# Stage B (Task expert + Router)
+lr_e2: 5.0e-4           
+
+text_lr_e2: 1.0e-5
+lr_e2_router: 1.0e-3
+lr_top_router: 1.0e-5
+lambda_b_con: 0.001
+tau_b_con: 0.19
+num_task_experts: 2
+seed: 43

configs/class/dtd.yaml

@@ -0,0 +1,79 @@
+hydra:
+  run:
+    dir: ./experiments/DTD
+  job:
+    chdir: true
+
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '%(message)s'
+
+class_order: ""
+dataset_root: ""
+workdir: ""
+log_path: "metrics.json"
+model_name: "ViT-B/16"
+prompt_template: "a bad photo of a {}."
+zero_shot_eval: true
+pre_task_zero_shot_eval: False
+zs_mtil_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]  # indices of MTIL datasets (0..24) to evaluate in zero-shot; excludes any selected as downstream train_dataset
+ 
+
+
+batch_size: 128
+increment: ${initial_increment}
+initial_increment: 10  
+scenario: "class"
+dataset: "cifar100"
+
+# method: "lwf"
+# lr: 7.1e-6
+# ls: 0.2
+# we: 
+# avg_freq:  
+# ref_dataset:
+# ref_sentences: 
+
+weight_decay: 0.0
+l2: 0
+ce_method: 0
+
+method: "MoE-Adapters"
+lr: 1e-3
+ls: 0.0
+#we:
+#avg_freq:
+#ref_dataset:
+#ref_sentences: random
+
+# =====================
+# DFA (two-expert + router) hyperparameters
+# =====================
+
+
+# Training schedule
+epochs_a: 1             
+epochs_b: 1             
+e2_top_k: 2             # top-k for E2 router  
+
+# Stage A (InfoNCE)
+lr_e1: 1.0e-5         
+tau_con: 0.15           # contrastive temperature for InfoNCE
+use_moco_queue: true    # enable momentum queue for more negatives
+moco_queue_size: 128   # queue size (default 4096, increases negative samples significantly)
+ 
+
+ 
+
+# Stage B (Task expert + Router)
+lr_e2: 5.0e-4           
+
+text_lr_e2: 1.0e-5
+lr_e2_router: 1.0e-3
+lr_top_router: 1.0e-5
+lambda_b_con: 0.001
+tau_b_con: 0.19
+num_task_experts: 2
+seed: 43