Function build_zero_shot_classifier

tools/test_zero_shot_hf.py:342–394

Build zero-shot classifier weights. Args: model: VTPModel instance with get_clip_text_feature method. tokenizer: Text tokenizer. classnames: Sequence of class names. templates: Sequence of template functions. num_classes_per_batch: Number of classes to pr

(
    model: VTPModel,
    tokenizer,
    classnames: Sequence[str],
    templates: Sequence[Callable],
    num_classes_per_batch: int = 10,
    device: torch.device = None,
    use_tqdm: bool = True,
)

# ============================================================================

def build_zero_shot_classifier(
    model: VTPModel,
    tokenizer,
    classnames: Sequence[str],
    templates: Sequence[Callable],
    num_classes_per_batch: int = 10,
    device: torch.device = None,
    use_tqdm: bool = True,
) -> torch.Tensor:
    """Build zero-shot classifier weights.

    Args:
        model: VTPModel instance with get_clip_text_feature method.
        tokenizer: Text tokenizer.
        classnames: Sequence of class names.
        templates: Sequence of template functions.
        num_classes_per_batch: Number of classes to process per batch.
        device: Device to use.
        use_tqdm: Whether to show progress bar.

    Returns:
        Zero-shot classifier weights, shape (embed_dim, num_classes).
    """
    num_templates = len(templates)
    num_classes = len(classnames)

    if use_tqdm:
        num_iter = (num_classes - 1) // num_classes_per_batch + 1
        iter_wrap = partial(tqdm, total=num_iter, desc="Building classifier", unit_scale=num_classes_per_batch)
    else:
        iter_wrap = iter

    def _process_batch(batch_classnames: List[str]) -> torch.Tensor:
        num_batch_classes = len(batch_classnames)
        texts = [template(c) for c in batch_classnames for template in templates]
        tokens = tokenizer(texts).to(device)

        # Use VTPModel's get_clip_text_feature method
        text_features = model.get_clip_text_feature(tokens, normalize=True)

        # Reshape and average over templates
        text_features = text_features.reshape(num_batch_classes, num_templates, -1).mean(dim=1)
        text_features = F.normalize(text_features, dim=1)
        return text_features.T  # (embed_dim, num_batch_classes)

    with torch.no_grad():
        batched_embeds = [
            _process_batch(batch)
            for batch in iter_wrap(batched(classnames, num_classes_per_batch))
        ]
        zeroshot_weights = torch.cat(batched_embeds, dim=1)

    return zeroshot_weights


# ============================================================================