main.py

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault(cuda_visible_devices := "CUDA_VISIBLE_DEVICES", "2")
import math
import json
import hydra
import logging
from omegaconf import DictConfig, ListConfig

from tqdm import tqdm

import torch
import numpy as np
import statistics
from torch.utils.data import DataLoader

import clip.clip as clip
from mtil_datasets import get_dataset as get_mtil_dataset
from continual_clip.clip_original import clip as clip_orig
from MTIL_datasets.voc2007 import VOC2007 as MTILVOC2007
from PIL import Image

from continual_clip import utils
from continual_clip.models import load_model
from continual_clip.datasets import build_cl_scenarios, get_dataset

MTIL_INDEX_TO_NAME = {
    0: "FGVCAircraft",
    1: "Caltech101",
    2: "CIFAR100",
    3: "DescribableTextures",
    4: "EuroSAT",
    5: "OxfordFlowers",
    6: "Food101",
    7: "MNIST",
    8: "OxfordPets",
    9: "StanfordCars",
    10: "SUN397",
    11: "Country211",
    12: "SST2",
    13: "HatefulMemes",
    14: "GTSRB",
    15: "RESISC45",
    16: "FER2013",
    17: "UCF101",
    18: "CIFAR10",
    19: "STL10",
    20: "VOC2007",
    21: "ImageNetR",
    22: "KittiDistance",
    23: "PCam",
    24: "CLEVRCount",
}

# Map MTIL indices to the dataset keys used by the downstream CIL loader.
TRAIN_INDEX_TO_DATASET_KEY = {
    0: "aircraft",
    1: "caltech101",
    2: "cifar100",
    3: "dtd",
    4: "eurosat",
    5: "oxford_flowers",
    6: "food101",
    7: "mnist",
    8: "oxford_pets",
    9: "stanford_cars",
    10: "sun397",
    11: "country211",
    12: "sst2",
    13: "hatefulmemes",
    14: "gtsrb",
    15: "resisc45",
    16: "fer2013",
    17: "ucf101",
    18: "cifar10",
    19: "stl10",
    20: "voc2007",
    21: "imagenet_r",
    22: "kitti_distance",
    23: "pcam",
    24: "clevr_count",
}

def evaluate_zero_shot(
    model, device, cfg, limit_datasets=None, use_original_clip=False
):
    """Evaluate zero-shot retention on MTIL auxiliary domains.

    The downstream training dataset is excluded so this routine measures the
    pre-trained knowledge retention side of the DFA-CIL protocol.
    """
    # Build a minimal config compatible with the MTIL dataset helper.
    class _ZSCfg:
        pass

    zs_cfg = _ZSCfg()
    zs_cfg.dataset = "MTIL"
    zs_cfg.dataset_root = cfg.dataset_root
    zs_cfg.seed = getattr(cfg, "seed", 1)
    zs_cfg.use_validation = getattr(cfg, "use_validation", False)
    zs_cfg.MTIL_order_2 = getattr(cfg, "MTIL_order_2", False)
    # Load the full MTIL pool first; the downstream training domain is filtered later.
    zs_cfg.train_one_dataset = -1

    # Choose between the adapted model and the original frozen CLIP baseline.
    orig_model = None
    tokenizer = clip.tokenize
    zs_transforms = getattr(model, "transforms", None)
    if use_original_clip:
        try:
            orig_model, _, zs_transforms = clip_orig.load(
                cfg.model_name, device=device, jit=False
            )
            orig_model.eval()
            tokenizer = clip_orig.tokenize
        except Exception as e:
            logging.error(f"Failed to load original CLIP for pre-task ZS: {e}")
            return {}

    try:
        zs_datasets, zs_classnames, zs_templates, zs_names = get_mtil_dataset(
            zs_cfg, split="test", transforms=zs_transforms
        )
    except Exception as e:
        logging.error(f"Zero-shot dataset loading failed: {e}")
        return {}

    # Optional allow-list for the retained zero-shot evaluation domains.
    zs_filter = getattr(cfg, "zero_shot_datasets", None)
    if isinstance(zs_filter, str):
        zs_filter = [s.strip() for s in zs_filter.split(",") if s.strip()]

    # Hydra configs may pass the MTIL allow-list in several container formats.
    def _parse_list(val):
        if isinstance(val, ListConfig):
            return [int(v) for v in val]
        if isinstance(val, (list, tuple)):
            return [int(v) for v in val]
        if isinstance(val, str):
            parts = [p.strip() for p in val.replace(";", ",").split(",") if p.strip()]
            return [int(p) for p in parts]
        if isinstance(val, (int,)):
            return [int(val)]
        return []

    zs_indices = _parse_list(getattr(cfg, "zs_mtil_indices", []))
    allowed_by_indices = {
        MTIL_INDEX_TO_NAME[i] for i in zs_indices if i in MTIL_INDEX_TO_NAME
    }

    max_zs_samples = int(getattr(cfg, "max_zs_samples", -1))  # -1 keeps the full dataset.
    zs_bs = int(getattr(cfg, "zs_batch_size", 32))
    num_workers = int(getattr(cfg, "num_workers", 4))
    pin_memory = device.type == "cuda"

    # Materialize the evaluation pool and apply the downstream / user filters.
    datasets_info = list(zip(zs_datasets, zs_classnames, zs_templates, zs_names))
    filtered = []

    # Exclude the downstream CIL dataset(s) from auxiliary zero-shot retention evaluation.
    train_indices = _parse_list(getattr(cfg, "train_dataset", []))
    # Backward compatibility for older single-dataset configs.
    if not train_indices:
        toi = int(getattr(cfg, "train_one_dataset", -1))
        if toi >= 0:
            train_indices = [toi]
    skip_names = {MTIL_INDEX_TO_NAME.get(i, "StanfordCars") for i in train_indices}
    # By default, evaluate all remaining MTIL domains once the downstream domain is removed.
    if not allowed_by_indices:
        all_names = {name for (_, _, _, name) in datasets_info}
        allowed_by_indices = all_names - skip_names
    for ds, classnames, templates, name in datasets_info:
        if name in skip_names:
            continue
        if zs_filter and name not in zs_filter:
            continue
        if allowed_by_indices and name not in allowed_by_indices:
            continue
        filtered.append((ds, classnames, templates, name))
    # Optional cap for exploratory runs after all domain filters have been applied.
    if isinstance(limit_datasets, int) and limit_datasets > 0:
        filtered = filtered[:limit_datasets]

    results = {}
    for ds, classnames, templates, name in filtered:
        # Use the dataset-provided zero-shot template when available.
        tmpl = None
        if isinstance(templates, (list, tuple)) and len(templates) > 0:
            tmpl = templates[0]

        def render(c):
            if callable(tmpl):
                try:
                    return tmpl(c)
                except Exception:
                    return f"a photo of a {c}."
            if isinstance(tmpl, str):
                try:
                    return tmpl.format(c)
                except Exception:
                    return f"a photo of a {c}."
            # Fall back to the run-level prompt template.
            try:
                return cfg.prompt_template.format(c)
            except Exception:
                return f"a photo of a {c}."

        prompts = [render(c) for c in classnames]
        try:
            text_tokens = tokenizer(prompts).to(device)
        except Exception as e:
            logging.error(
                f"Tokenization failed for {name}: {e}. Prompts sample: "
                f"{prompts[:3] if len(prompts) > 3 else prompts}"
            )
            continue

        # VOC2007 stays multi-label for mAP; all other domains are reduced to single labels.
        def _zs_collate(batch):
            xs = []
            ys = []
            if name == "VOC2007":
                for xi, yi in batch:
                    xs.append(xi)
                    if isinstance(yi, torch.Tensor):
                        yv = yi.detach().cpu().numpy()
                    elif isinstance(yi, (list, tuple, np.ndarray)):
                        yv = np.asarray(yi)
                    else:
                        # Robust fallback for unexpected scalar labels.
                        vec = np.zeros(len(classnames), dtype=np.int64)
                        try:
                            vec[int(yi)] = 1
                        except Exception:
                            pass
                        yv = vec
                    yv = np.asarray(yv).astype(np.int64).reshape(-1)
                    if len(yv) != len(classnames):
                        # Coerce malformed vectors back to the expected one-hot length.
                        vec = np.zeros(len(classnames), dtype=np.int64)
                        try:
                            vec[int(np.argmax(yv))] = 1
                        except Exception:
                            pass
                        yv = vec
                    ys.append(torch.tensor(yv, dtype=torch.long))
                x_batch = torch.stack(xs, dim=0)
                y_batch = torch.stack(ys, dim=0)  # [B, C] multi-hot labels.
                return x_batch, y_batch
            else:
                for xi, yi in batch:
                    xs.append(xi)
                    # Collapse vector-like labels to a scalar class id.
                    if isinstance(yi, torch.Tensor):
                        arr = yi.detach().cpu().numpy()
                    elif isinstance(yi, (list, tuple, np.ndarray)):
                        arr = np.asarray(yi)
                    else:
                        arr = yi
                    if isinstance(arr, (list, tuple, np.ndarray)):
                        arr = np.asarray(arr)
                        if arr.ndim == 0:
                            yi_scalar = int(arr.item())
                        else:
                            yi_scalar = int(arr.argmax())
                    else:
                        yi_scalar = int(arr)
                    ys.append(yi_scalar)
                x_batch = torch.stack(xs, dim=0)
                y_batch = torch.tensor(ys, dtype=torch.long)
                return x_batch, y_batch

        loader = DataLoader(
            ds,
            batch_size=zs_bs,
            num_workers=num_workers,
            pin_memory=pin_memory,
            collate_fn=_zs_collate,
        )
        correct = 0
        total = 0
        processed = 0
        # VOC2007 is reported with 11-point mAP instead of top-1 accuracy.
        voc_y_true = []
        voc_y_score = []
        with torch.inference_mode():
            # When measuring A_k^0, reuse frozen original CLIP text features across the dataset.
            if use_original_clip and orig_model is not None:
                text_features = orig_model.encode_text(text_tokens)
                text_features = text_features / text_features.norm(dim=-1, keepdim=True)

            for x, y in tqdm(loader, desc=f"ZS {name}", leave=False):
                x = x.to(device, non_blocking=True)
                # Preserve multi-label targets for VOC2007; otherwise build a 1D class tensor.
                if name == "VOC2007":
                    # Ensure shape [B, C].
                    if isinstance(y, torch.Tensor):
                        y_vec = y
                    else:
                        y_vec = torch.as_tensor(y)
                    if y_vec.ndim == 1 and y_vec.numel() == len(classnames):
                        y_vec = y_vec.view(1, -1)
                else:
                    # Robust single-label conversion for heterogeneous dataset wrappers.
                    def _to_label_tensor(y_any):
                        if isinstance(y_any, torch.Tensor):
                            if y_any.ndim > 1:
                                y_any = y_any.argmax(dim=1)
                            return y_any.to(device, non_blocking=True).long()
                        if isinstance(y_any, (list, tuple)):
                            proc = []
                            for elem in y_any:
                                if isinstance(elem, torch.Tensor):
                                    if elem.ndim == 0:
                                        proc.append(int(elem.item()))
                                    else:
                                        proc.append(
                                            int(elem.detach().cpu().numpy().argmax())
                                        )
                                elif isinstance(elem, (list, tuple, np.ndarray)):
                                    arr = np.asarray(elem)
                                    if arr.ndim == 0:
                                        proc.append(int(arr.item()))
                                    else:
                                        proc.append(int(arr.argmax()))
                                else:
                                    proc.append(int(elem))
                            return torch.tensor(proc, device=device, dtype=torch.long)
                        try:
                            return torch.tensor(
                                [int(y_any)], device=device, dtype=torch.long
                            )
                        except Exception:
                            return torch.tensor(y_any, device=device, dtype=torch.long)

                    y = _to_label_tensor(y)
                    bsz_now = x.size(0)
                    if y.ndim == 1 and y.size(0) != bsz_now:
                        if y.size(0) == len(classnames):
                            y = y.argmax(dim=0).reshape(1).to(device).long()
                        elif (y.numel() % max(1, len(classnames))) == 0 and len(
                            classnames
                        ) > 0:
                            try:
                                y = (
                                    y.view(-1, len(classnames))
                                    .argmax(dim=1)
                                    .to(device)
                                    .long()
                                )
                            except Exception:
                                pass
                        if y.ndim == 1 and y.size(0) != bsz_now:
                            if y.numel() == 1:
                                y = y.view(1).repeat(bsz_now).to(device)
                            else:
                                y = y[:bsz_now].to(device)
                    if not (
                        isinstance(y, torch.Tensor)
                        and y.ndim == 1
                        and y.size(0) == bsz_now
                    ):
                        y = torch.as_tensor(y, device=device)
                        y = y.view(-1)
                        if len(classnames) > 0 and y.numel() == len(classnames):
                            y = y.argmax().view(1).repeat(bsz_now)
                        elif len(classnames) > 0 and y.numel() == bsz_now * len(
                            classnames
                        ):
                            y = y.view(bsz_now, len(classnames)).argmax(dim=1)
                        elif y.numel() == 1:
                            y = y.view(1).repeat(bsz_now)
                        elif y.numel() > bsz_now:
                            y = y[:bsz_now]
                        else:
                            pad_val = int(y[0].item()) if y.numel() > 0 else 0
                            y = torch.nn.functional.pad(
                                y.long(), (0, bsz_now - y.numel()), value=pad_val
                            )
                        y = y.long()

                if use_original_clip and orig_model is not None:
                    image_features = orig_model.encode_image(x)
                    image_features = image_features / image_features.norm(
                        dim=-1, keepdim=True
                    )
                    logit_scale = getattr(orig_model, "logit_scale", None)
                    if logit_scale is not None and hasattr(logit_scale, "exp"):
                        scale = logit_scale.exp()
                    else:
                        scale = 1.0
                    logits = scale * image_features @ text_features.t()
                else:
                    # Use the unified DFA-MoE forward path when exposed by the model wrapper.
                    if hasattr(model, "compute_logits") and callable(
                        getattr(model, "compute_logits")
                    ):
                        logits = model.compute_logits(x, text_tokens)
                    else:
                        logits, _ = model.model(x, text_tokens, 0, is_train=False)

                if name == "VOC2007":
                    # Accumulate predictions for VOC2007 mAP computation.
                    voc_y_score.append(logits.detach().cpu())
                    voc_y_true.append(y_vec.detach().cpu())
                    processed += x.size(0)
                    if max_zs_samples > 0 and processed >= max_zs_samples:
                        break
                    continue

                pred = logits.argmax(dim=1)
                correct += (pred == y).sum().item()
                bsz = y.size(0)
                total += bsz
                processed += bsz
                if max_zs_samples > 0 and processed >= max_zs_samples:
                    break
        # Release per-domain tensors before moving to the next auxiliary dataset.
        del text_tokens
        if use_original_clip and orig_model is not None:
            try:
                del text_features
            except Exception:
                pass
        torch.cuda.empty_cache()
        if name == "VOC2007":
            # Compute the standard VOC2007 11-point mAP.
            def _ap11(y_true_cls: np.ndarray, y_score_cls: np.ndarray) -> float:
                # Rank examples by descending confidence.
                order = np.argsort(-y_score_cls)
                y_true_sorted = y_true_cls[order]
                tp = (y_true_sorted == 1).astype(np.float32)
                fp = (y_true_sorted == 0).astype(np.float32)
                tp_cum = np.cumsum(tp)
                fp_cum = np.cumsum(fp)
                # Numerical safeguard for empty precision denominators.
                prec = tp_cum / np.maximum(tp_cum + fp_cum, 1e-12)
                # Recall normalized by the number of positives for the class.
                total_pos = max(1.0, float((y_true_cls == 1).sum()))
                rec = tp_cum / total_pos
                ap = 0.0
                for r in np.linspace(0.0, 1.0, 11):
                    mask = rec >= r
                    p_interp = np.max(prec[mask]) if np.any(mask) else 0.0
                    ap += p_interp
                return ap / 11.0

            if voc_y_true and voc_y_score:
                y_true_all = torch.cat(voc_y_true, dim=0).numpy()
                y_score_all = torch.cat(voc_y_score, dim=0).numpy()
                aps = []
                for ci in range(y_true_all.shape[1]):
                    aps.append(
                        _ap11(
                            y_true_all[:, ci].astype(np.int64),
                            y_score_all[:, ci].astype(np.float32),
                        )
                    )
                mAP = float(np.mean(aps)) if aps else 0.0
                results[name] = round(100.0 * mAP, 2)
            else:
                results[name] = 0.0
        else:
            acc = 100.0 * correct / total if total > 0 else 0.0
            results[name] = round(acc, 2)

    return results


class TaskIdOffsetDataset(torch.utils.data.Dataset):
    """Replace local task ids with global incremental-task ids for evaluation bookkeeping."""

    def __init__(self, ds, offset: int):
        self.ds = ds
        self.offset = int(offset)

    def __len__(self):
        return len(self.ds)

    def __getitem__(self, idx):
        x, y, t = self.ds[idx]
        # The wrapped scenario already defines the sample; only the task id is remapped.
        return x, y, int(self.offset)


def _parse_int_list(val):
    if isinstance(val, ListConfig):
        return [int(v) for v in val]
    if isinstance(val, (list, tuple)):
        return [int(v) for v in val]
    if isinstance(val, str):
        parts = [p.strip() for p in val.replace(";", ",").split(",") if p.strip()]
        return [int(p) for p in parts]
    if isinstance(val, (int,)):
        return [int(val)]
    return []


@hydra.main(config_path=None, config_name=None, version_base="1.1")
def continual_clip(cfg: DictConfig) -> None:

    cfg.workdir = utils.get_workdir(path=os.getcwd())
    # Resolve relative dataset paths from the Hydra work directory.
    try:
        if not os.path.isabs(str(getattr(cfg, "dataset_root", ""))):
            cfg.dataset_root = os.path.join(cfg.workdir, cfg.dataset_root)
    except Exception:
        cfg.dataset_root = os.path.join(cfg.workdir, cfg.dataset_root)

    # Seed all RNGs so CIL order, queue sampling, and evaluation are reproducible.
    utils.seed_all(int(getattr(cfg, "seed", 1)))

    train_indices = _parse_int_list(getattr(cfg, "train_dataset", []))
    if not train_indices:
        train_one_dataset = int(getattr(cfg, "train_one_dataset", -1))
        if train_one_dataset >= 0:
            train_indices = [train_one_dataset]
    if not train_indices:
        raise ValueError("Please provide a single train_dataset index (0..24).")
    if len(train_indices) != 1:
        raise ValueError(
            f"Only a single downstream dataset is supported. Got train_dataset={train_indices}"
        )

    splits_list = _parse_int_list(getattr(cfg, "cil_splits", []))
    if not splits_list:
        raise ValueError("Please provide a single cil_splits value.")
    if len(splits_list) != 1:
        raise ValueError(
            f"Only a single cil_splits value is supported. Got cil_splits={splits_list}"
        )

    train_index = int(train_indices[0])
    cil_splits = int(splits_list[0])
    cfg.train_dataset = train_index
    cfg.cil_splits = cil_splits

    utils.save_config(cfg)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # External class orders are optional; otherwise the scenario defines the split.
    if getattr(cfg, "class_order", None):
        cfg.class_order = utils.get_class_order(
            os.path.join(cfg.workdir, cfg.class_order)
        )
    else:
        cfg.class_order = None
    model = load_model(cfg, device)

    if train_index not in TRAIN_INDEX_TO_DATASET_KEY:
        raise ValueError(
            f"train_dataset contains invalid index {train_index}. Supported indices: {list(TRAIN_INDEX_TO_DATASET_KEY.keys())}"
        )
    dataset_key = TRAIN_INDEX_TO_DATASET_KEY[train_index]
    cfg.dataset = dataset_key
    try:
        _, _tmp_classes = get_dataset(cfg, is_train=True)
        num_classes = len(_tmp_classes)
    except Exception:
        fallback_classes = {
            "cifar100": 100,
            "stanford_cars": 196,
        }
        num_classes = fallback_classes.get(cfg.dataset, 100)
    inc = math.ceil(num_classes / cil_splits)
    cfg.initial_increment = inc
    cfg.increment = inc
    cfg.cil_splits = cil_splits

    eval_scenario, _ = build_cl_scenarios(
        cfg, is_train=False, transforms=model.transforms
    )
    train_scenario, train_classes = build_cl_scenarios(
        cfg, is_train=True, transforms=model.transforms
    )

    try:

        class _TCfg:
            pass

        tcfg = _TCfg()
        tcfg.dataset = "MTIL"
        tcfg.dataset_root = cfg.dataset_root
        tcfg.seed = getattr(cfg, "seed", 1)
        tcfg.use_validation = getattr(cfg, "use_validation", False)
        tcfg.MTIL_order_2 = getattr(cfg, "MTIL_order_2", False)
        tcfg.train_one_dataset = train_index
        _, _, _tmpl_tmp, _ = get_mtil_dataset(
            tcfg, split="test", transforms=model.transforms
        )
        templates_first = None
        templates_list = None
        if isinstance(_tmpl_tmp, (list, tuple)) and len(_tmpl_tmp) > 0:
            per_ds_templates = _tmpl_tmp[0]
            if isinstance(per_ds_templates, (list, tuple)) and len(per_ds_templates) > 0:
                templates_list = list(per_ds_templates)
                templates_first = per_ds_templates[0]
            elif isinstance(per_ds_templates, str) or callable(per_ds_templates):
                templates_list = [per_ds_templates]
                templates_first = per_ds_templates
    except Exception:
        templates_first = None
        templates_list = None

    with open(cfg.log_path, "w+") as f:
        pass

    acc_list = []
    # Accuracy recorded when each incremental task is first learned, used for BWT.
    acc_at_learn_time = {}
    # Preserve the absolute class ids introduced by each incremental task.
    block_abs_ids = []

    # Optional A_k^0 baseline: evaluate the original frozen CLIP before any adaptation.
    if bool(getattr(cfg, "pre_task_zero_shot_eval", True)) and bool(
        getattr(cfg, "zero_shot_eval", True)
    ):
        skip_name = MTIL_INDEX_TO_NAME.get(train_index, "StanfordCars")
        logging.info(
            f"Pre-task zero-shot evaluation with ORIGINAL CLIP (excluding {skip_name})..."
        )
        zs_pre_results = evaluate_zero_shot(model, device, cfg, use_original_clip=True)
        with open(cfg.log_path, "a+") as f:
            f.write(
                json.dumps(
                    {
                        "task": -1,
                        "zs_pre": zs_pre_results,
                    }
                )
                + "\n"
            )

    # Standard class-incremental training and evaluation on one downstream dataset.
    for task_id in range(len(train_scenario)):
        logging.info(f"Evaluation for task {task_id} has started.")
        model.classes_names = train_classes
        cfg.initial_increment = inc
        cfg.increment = inc
        model.class_ids_per_task = None
        model.adaptation(task_id, cfg, train_scenario, train_classes)
        # Record which absolute classes became visible at this incremental step.
        abs_ids = list(getattr(model, "last_task_real_ids", []))
        block_abs_ids.append(abs_ids)
        # Clear allocator state before the evaluation phase.
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        # Evaluate strict CIL over the cumulative seen-class label space.
        eval_bs = int(getattr(cfg, "eval_batch_size", 32))
        text_chunk = int(getattr(cfg, "eval_text_chunk", 512))
        global_seen = []
        for g_idx_seen in range(task_id + 1):
            for cid in block_abs_ids[g_idx_seen]:
                global_seen.append(int(cid))

        def _render_with_template(class_name):
            tmpl = templates_first
            if callable(tmpl):
                try:
                    return tmpl(class_name)
                except Exception:
                    pass
            if isinstance(tmpl, str):
                try:
                    return tmpl.format(class_name)
                except Exception:
                    pass
            # Fall back to the run-level prompt template.
            try:
                return cfg.prompt_template.format(class_name)
            except Exception:
                return f"a photo of a {class_name}."

        prompts_all = []
        token_to_class_index = []
        class_template_counts = [0 for _ in range(len(global_seen))]
        for g_idx, cid in enumerate(global_seen):
            name = train_classes[cid]
            tlist = templates_list if templates_list else None
            if not tlist:
                prompts_all.append(_render_with_template(name))
                token_to_class_index.append(g_idx)
                class_template_counts[g_idx] += 1
            else:
                for t in tlist:
                    if callable(t):
                        try:
                            s = t(name)
                        except Exception:
                            s = _render_with_template(name)
                    elif isinstance(t, str):
                        try:
                            s = t.format(name)
                        except Exception:
                            s = _render_with_template(name)
                    else:
                        s = _render_with_template(name)
                    prompts_all.append(s)
                    token_to_class_index.append(g_idx)
                    class_template_counts[g_idx] += 1
        tokens_all = clip.tokenize(
            prompts_all
        )  # Keep on CPU and stream prompt chunks to the device on demand.
        global_index_of = {cid: i for i, cid in enumerate(global_seen)}
        task_correct = {}
        task_total = {}
        for g_idx in range(task_id + 1):
            ds = eval_scenario[g_idx]
            loader = DataLoader(
                TaskIdOffsetDataset(ds, offset=g_idx), batch_size=eval_bs
            )
            with torch.no_grad():
                # Average logits over the number of templates assigned to each class.
                counts_tensor = torch.tensor(
                    class_template_counts, dtype=torch.float32, device=device
                ).clamp_min(1.0)
                for inputs, targets, _task_ids in tqdm(loader):
                    inputs = inputs.to(device, non_blocking=True)
                    batch_size = inputs.shape[0]
                    # Aggregate per-template logits into a single score per seen class.
                    agg_logits = torch.zeros(
                        (batch_size, len(global_seen)), device=device
                    )
                    for start in range(0, tokens_all.size(0), max(1, text_chunk)):
                        end = min(tokens_all.size(0), start + max(1, text_chunk))
                        chunk = tokens_all[start:end].to(device, non_blocking=True)
                        if hasattr(model, "compute_logits") and callable(
                            getattr(model, "compute_logits")
                        ):
                            logits_chunk = model.compute_logits(inputs, chunk)
                        else:
                            logits_chunk, _ = model.model(
                                inputs, chunk, 0, is_train=False
                            )
                        # Scatter template logits back to their corresponding class slot.
                        idx_chunk = (
                            torch.tensor(
                                token_to_class_index[start:end], device=device
                            )
                            .view(1, -1)
                            .expand(batch_size, -1)
                        )
                        if logits_chunk.dtype != agg_logits.dtype:
                            logits_chunk = logits_chunk.to(dtype=agg_logits.dtype)
                        agg_logits.scatter_add_(1, idx_chunk, logits_chunk)
                    # Convert summed template logits into mean class logits.
                    agg_logits = agg_logits / counts_tensor.view(1, -1)
                    preds_global = agg_logits.detach().cpu().argmax(dim=1).numpy()
                    # Convert absolute dataset labels into indices of the seen-class bank.
                    if isinstance(targets, torch.Tensor):
                        t_np = targets.detach().cpu().numpy()
                    else:
                        t_np = np.asarray(targets)
                    mapped = np.array(
                        [global_index_of.get(int(v), -1) for v in t_np],
                        dtype=np.int64,
                    )
                    valid = mapped >= 0
                    corr = int((preds_global[valid] == mapped[valid]).sum())
                    tot = int(valid.sum())
                    task_correct[g_idx] = task_correct.get(g_idx, 0) + corr
                    task_total[g_idx] = task_total.get(g_idx, 0) + tot
        # Release the per-step text bank before the next task.
        del tokens_all
        torch.cuda.empty_cache()

        # VOC2007 remains multi-label, so CIL is reported with mAP instead of top-1 accuracy.
        voc_mAP = None
        voc_tid_override = None
        try:
            voc_tids = list(range(task_id + 1)) if dataset_key == "voc2007" else []
            if voc_tids:
                tlist_voc = templates_list if templates_list else None

                def _render_voc_all(cname: str):
                    outs = []
                    if tlist_voc:
                        for t in tlist_voc:
                            if callable(t):
                                try:
                                    outs.append(t(cname))
                                except Exception:
                                    continue
                            elif isinstance(t, str):
                                try:
                                    outs.append(t.format(cname))
                                except Exception:
                                    continue
                    if not outs:
                        try:
                            outs = [cfg.prompt_template.format(cname)]
                        except Exception:
                            outs = [f"a photo of a {cname}."]
                    return outs

                # Build the full prompt bank for the 20 VOC classes.
                voc_ds_multi = MTILVOC2007(
                    root=cfg.dataset_root,
                    seed=getattr(cfg, "seed", 1),
                    single_label=False,
                )
                voc_prompts = []
                voc_token_to_class = []
                for ci, cname in enumerate(voc_ds_multi.classnames):
                    outs = _render_voc_all(cname)
                    voc_prompts.extend(outs)
                    voc_token_to_class.extend([ci] * len(outs))
                voc_tokens = clip.tokenize(voc_prompts).to(device)
                # Count templates per class so logits can be averaged back to class level.
                voc_counts = torch.zeros(
                    len(voc_ds_multi.classnames), dtype=torch.float32, device=device
                )
                for ci in voc_token_to_class:
                    voc_counts[ci] += 1.0
                # Stream the VOC test set in mini-batches.
                y_true = []
                y_score = []
                batch = []

                def _flush_batch(batch_list):
                    if not batch_list:
                        return
                    imgs = []
                    ys = []
                    for d in batch_list:
                        try:
                            img = Image.open(d.impath).convert("RGB")
                            if getattr(model, "transforms", None) is not None:
                                img = model.transforms(img)
                            imgs.append(img)
                            ys.append(torch.tensor(d.label, dtype=torch.long))
                        except Exception:
                            continue
                    if not imgs:
                        return
                    x = torch.stack(imgs, dim=0).to(device, non_blocking=True)
                    with torch.no_grad():
                        if hasattr(model, "compute_logits") and callable(
                            getattr(model, "compute_logits")
                        ):
                            logits_full = model.compute_logits(x, voc_tokens)
                        else:
                            logits_full, _ = model.model(
                                x, voc_tokens, 0, is_train=False
                            )
                    # Collapse prompt-level logits back to class-level logits.
                    B = logits_full.size(0)
                    Gv = len(voc_ds_multi.classnames)
                    agg = torch.zeros((B, Gv), device=logits_full.device)
                    idx_chunk = (
                        torch.tensor(voc_token_to_class, device=logits_full.device)
                        .view(1, -1)
                        .expand(B, -1)
                    )
                    if logits_full.dtype != agg.dtype:
                        logits_full = logits_full.to(dtype=agg.dtype)
                    agg.scatter_add_(1, idx_chunk, logits_full)
                    agg = agg / voc_counts.view(1, -1)
                    y_score.append(agg.detach().cpu())
                    y_true.append(torch.stack(ys, dim=0))

                bs_local = eval_bs
                for d in voc_ds_multi.test:
                    batch.append(d)
                    if len(batch) >= bs_local:
                        _flush_batch(batch)
                        batch = []
                if batch:
                    _flush_batch(batch)
                if y_true and y_score:
                    y_true_all = torch.cat(y_true, dim=0).numpy()
                    y_score_all = torch.cat(y_score, dim=0).numpy()

                    # Per-class 11-point AP, then macro-average over classes.
                    def _ap11(y_true_cls: np.ndarray, y_score_cls: np.ndarray) -> float:
                        order = np.argsort(-y_score_cls)
                        y_true_sorted = y_true_cls[order]
                        tp = (y_true_sorted == 1).astype(np.float32)
                        fp = (y_true_sorted == 0).astype(np.float32)
                        tp_cum = np.cumsum(tp)
                        fp_cum = np.cumsum(fp)
                        prec = tp_cum / np.maximum(tp_cum + fp_cum, 1e-12)
                        total_pos = max(1.0, float((y_true_cls == 1).sum()))
                        rec = tp_cum / total_pos
                        ap = 0.0
                        for r in np.linspace(0.0, 1.0, 11):
                            mask = rec >= r
                            p_interp = np.max(prec[mask]) if np.any(mask) else 0.0
                            ap += p_interp
                        return ap / 11.0

                    aps = []
                    for ci in range(y_true_all.shape[1]):
                        aps.append(
                            _ap11(
                                y_true_all[:, ci].astype(np.int64),
                                y_score_all[:, ci].astype(np.float32),
                            )
                        )
                    voc_mAP = 100.0 * float(np.mean(aps)) if aps else None
                    voc_tid_override = voc_tids[-1]
        except Exception as e:
            logging.error(f"VOC2007 mAP (CIL) failed: {e}")

        # Auxiliary zero-shot retention evaluation for PKF tracking.
        zs_results = {}
        if getattr(cfg, "zero_shot_eval", True):
            zs_results = evaluate_zero_shot(model, device, cfg)
            # Convenience summary; SCR is computed later from the raw per-domain scores.
            if zs_results:
                zs_mean = round(sum(zs_results.values()) / len(zs_results), 2)
            else:
                zs_mean = 0.0

        # Aggregate CIL metrics over all seen tasks.
        seen_task_ids = list(range(task_id + 1))
        acc_per_task = []
        for tid in seen_task_ids:
            tot = task_total.get(tid, 0)
            if (
                voc_tid_override is not None
                and tid == voc_tid_override
                and voc_mAP is not None
            ):
                acc = max(0.0, min(1.0, voc_mAP / 100.0))
            else:
                acc = (task_correct.get(tid, 0) / tot) if tot > 0 else 0.0
            acc_per_task.append(acc)
        # Overall CIL score: replace the VOC task contribution with its mAP estimate.
        if (
            voc_tid_override is not None
            and voc_mAP is not None
            and voc_tid_override in task_total
        ):
            total_samples = max(1, sum(task_total.values()))
            corrected_sum = 0.0
            for tid in seen_task_ids:
                if tid == voc_tid_override:
                    corrected_sum += (voc_mAP / 100.0) * task_total.get(tid, 0)
                else:
                    corrected_sum += task_correct.get(tid, 0)
            overall_acc = 100.0 * (corrected_sum / total_samples)
        else:
            overall_acc = 100.0 * (
                sum(task_correct.values()) / max(1, sum(task_total.values()))
            )
        acc_list.append(overall_acc)
        # Diagonal entry in the CIL accuracy matrix, used as the BWT reference.
        if task_id not in acc_at_learn_time:
            acc_at_learn_time[task_id] = acc_per_task[task_id]
        # BWT follows the standard mean difference from the learn-time accuracy.
        bwt_vals = []
        for tid in seen_task_ids[:-1]:
            base = acc_at_learn_time.get(tid, acc_per_task[tid])
            bwt_vals.append(acc_per_task[tid] - base)
        bwt_val = round(100.0 * (sum(bwt_vals) / max(1, len(bwt_vals))), 2)
        acc_per_task_list = [round(100.0 * a, 2) for a in acc_per_task]
        with open(cfg.log_path, "a+") as f:
            f.write(
                json.dumps(
                    {
                        "task": task_id,
                        "acc": round(overall_acc, 2),
                        "acc_per_task": acc_per_task_list,
                        "bwt": bwt_val,
                        "zs": zs_results,
                    }
                )
                + "\n"
            )

        # Persist a compact zero-shot summary alongside the full metric log.
        if getattr(cfg, "zero_shot_eval", True) and zs_results:
            zs_mean_path = cfg.log_path.replace(".json", "_zs_mean.json")
            with open(zs_mean_path, "a+") as f:
                f.write(
                    json.dumps(
                        {
                            "task": task_id,
                            "zs_mean": zs_mean,
                            "zs_details": zs_results,
                        }
                    )
                    + "\n"
                )
    with open(cfg.log_path, "a+") as f:
        f.write(
            json.dumps(
                {
                    "last": round(acc_list[-1], 2),
                    "avg": round(statistics.mean(acc_list), 2),
                }
            )
            + "\n"
        )


if __name__ == "__main__":
    continual_clip()