birder-clip 0.0.2.dev7__tar.gz → 0.0.2.dev8__tar.gz

{birder_clip-0.0.2.dev7 → birder_clip-0.0.2.dev8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: birder_clip
-Version: 0.0.2.dev7
+Version: 0.0.2.dev8
 Summary: A Birder extension for CLIP-style image-text modeling and multimodal computer vision workflows.
 Author: Ofer Hasson
 License-Expression: Apache-2.0
@@ -24,7 +24,7 @@ Classifier: Typing :: Typed
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: birder>=0.6.0
+Requires-Dist: birder>=0.6.2
 Requires-Dist: ftfy>=6.3.1
 Requires-Dist: regex>=2025.7.29
 Requires-Dist: tqdm>=4.67.0
@@ -38,7 +38,7 @@ Requires-Dist: bandit~=1.9.4; extra == "dev"
 Requires-Dist: black~=26.5.0; extra == "dev"
 Requires-Dist: build~=1.5.0; extra == "dev"
 Requires-Dist: bumpver~=2026.1132; extra == "dev"
-Requires-Dist: coverage~=7.14.2; extra == "dev"
+Requires-Dist: coverage~=7.14.3; extra == "dev"
 Requires-Dist: debugpy; extra == "dev"
 Requires-Dist: flake8-pep585~=0.1.7; extra == "dev"
 Requires-Dist: flake8~=7.3.0; extra == "dev"
@@ -85,7 +85,7 @@ Full training is supported, but for large-scale CLIP pretraining you are probabl
 1. Ensure your environment meets the minimum requirements:
    - Python 3.11 or newer
    - PyTorch 2.10 or newer (installed for your hardware/driver stack)
-   - Birder 0.6.0 or newer
+   - Birder 0.6.2 or newer
 
 1. Install the latest Birder CLIP version:
 

{birder_clip-0.0.2.dev7 → birder_clip-0.0.2.dev8}/README.md

@@ -23,7 +23,7 @@ Full training is supported, but for large-scale CLIP pretraining you are probabl
 1. Ensure your environment meets the minimum requirements:
    - Python 3.11 or newer
    - PyTorch 2.10 or newer (installed for your hardware/driver stack)
-   - Birder 0.6.0 or newer
+   - Birder 0.6.2 or newer
 
 1. Install the latest Birder CLIP version:
 

{birder_clip-0.0.2.dev7 → birder_clip-0.0.2.dev8}/birder_clip/common/fs_ops.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from typing import Any
 from typing import NamedTuple
 from typing import Optional
+from typing import TypeAlias
 
 import torch
 from birder.common import cli
@@ -16,8 +17,10 @@ from birder.data.transforms.classification import inference_preset
 from birder_clip.common import lib
 from birder_clip.model_registry import Task
 from birder_clip.model_registry import registry
-from birder_clip.model_registry.manifest import EncoderMetadataType
 from birder_clip.model_registry.manifest import FileFormatType
+from birder_clip.model_registry.manifest import ImageEncoderMetadataType
+from birder_clip.model_registry.manifest import TextDecoderMetadataType
+from birder_clip.model_registry.manifest import TextEncoderMetadataType
 from birder_clip.net.base import BaseNet
 from birder_clip.net.base import SignatureType
 from birder_clip.tokenizers import Tokenizer
@@ -36,6 +39,8 @@ except ImportError:
 
 logger = logging.getLogger(__name__)
 
+ComponentMetadataType: TypeAlias = ImageEncoderMetadataType | TextEncoderMetadataType | TextDecoderMetadataType
+
 
 class ModelInfo(NamedTuple):
     signature: SignatureType
@@ -51,16 +56,18 @@ def write_config(network_name: str, net: BaseNet, signature: SignatureType, rgb_
         json.dump(model_config, handle, indent=2)
 
 
-def _split_encoder_metadata(encoder: Optional[EncoderMetadataType]) -> tuple[Optional[str], Optional[dict[str, Any]]]:
-    if encoder is None:
+def _split_component_metadata(
+    component: Optional[ComponentMetadataType],
+) -> tuple[Optional[str], Optional[dict[str, Any]]]:
+    if component is None:
         return (None, None)
-    if isinstance(encoder, str):
-        return (encoder, None)
+    if isinstance(component, str):
+        return (component, None)
 
-    if "network" not in encoder:
-        raise ValueError("Encoder metadata must include a 'network' field")
+    if "network" not in component:
+        raise ValueError("Component metadata must include a 'network' field")
 
-    return (None, encoder)  # type: ignore[return-value]
+    return (None, component)  # type: ignore[return-value]
 
 
 def model_path(
@@ -286,10 +293,12 @@ def load_checkpoint(
     tag: Optional[str] = None,
     image_encoder: Optional[str] = None,
     text_encoder: Optional[str] = None,
+    text_decoder: Optional[str] = None,
     embed_dim: Optional[int] = None,
     tokenizer: Optional[str] = None,
     image_encoder_config: Optional[dict[str, Any]] = None,
     text_encoder_config: Optional[dict[str, Any]] = None,
+    text_decoder_config: Optional[dict[str, Any]] = None,
     epoch: Optional[int] = None,
     new_size: Optional[tuple[int, int]] = None,
     new_context_length: Optional[int] = None,
@@ -300,6 +309,7 @@ def load_checkpoint(
         tag=tag,
         image_encoder=image_encoder,
         text_encoder=text_encoder,
+        text_decoder=text_decoder,
         embed_dim=embed_dim,
         tokenizer=tokenizer,
     )
@@ -329,10 +339,12 @@ def load_checkpoint(
         checkpoint_config,
         image_encoder=image_encoder,
         text_encoder=text_encoder,
+        text_decoder=text_decoder,
         embed_dim=embed_dim,
         tokenizer=tokenizer,
         image_encoder_config=image_encoder_config,
         text_encoder_config=text_encoder_config,
+        text_decoder_config=text_decoder_config,
         input_channels=input_channels,
         image_size=size,
         context_length=context_length,
@@ -364,10 +376,12 @@ def load_model(
     tag: Optional[str] = None,
     image_encoder: Optional[str] = None,
     text_encoder: Optional[str] = None,
+    text_decoder: Optional[str] = None,
     embed_dim: Optional[int] = None,
     tokenizer: Optional[str] = None,
     image_encoder_config: Optional[dict[str, Any]] = None,
     text_encoder_config: Optional[dict[str, Any]] = None,
+    text_decoder_config: Optional[dict[str, Any]] = None,
     epoch: Optional[int] = None,
     new_size: Optional[tuple[int, int]] = None,
     new_context_length: Optional[int] = None,
@@ -381,6 +395,7 @@ def load_model(
             tag=tag,
             image_encoder=image_encoder,
             text_encoder=text_encoder,
+            text_decoder=text_decoder,
             embed_dim=embed_dim,
             tokenizer=tokenizer,
         )
@@ -420,10 +435,12 @@ def load_model(
         checkpoint_config,
         image_encoder=image_encoder,
         text_encoder=text_encoder,
+        text_decoder=text_decoder,
         embed_dim=embed_dim,
         tokenizer=tokenizer,
         image_encoder_config=image_encoder_config,
         text_encoder_config=text_encoder_config,
+        text_decoder_config=text_decoder_config,
         input_channels=input_channels,
         image_size=size,
         context_length=context_length,
@@ -525,15 +542,17 @@ def load_pretrained_model(
     if model_metadata["task"] != Task.IMAGE_TEXT:
         raise ValueError(f"Unknown model type: {model_metadata['task']}")
 
-    image_encoder, image_config = _split_encoder_metadata(model_metadata["net"].get("image_encoder", None))
-    text_encoder, text_config = _split_encoder_metadata(model_metadata["net"].get("text_encoder", None))
+    image_encoder, image_config = _split_component_metadata(model_metadata["net"].get("image_encoder", None))
+    text_encoder, text_config = _split_component_metadata(model_metadata["net"].get("text_encoder", None))
+    text_decoder, decoder_config = _split_component_metadata(model_metadata["net"].get("text_decoder", None))
 
     pretrained_config: dict[str, Any] = {}
     if image_config is not None:
         pretrained_config["image"] = image_config
-
     if text_config is not None:
         pretrained_config["text"] = text_config
+    if decoder_config is not None:
+        pretrained_config["decoder"] = decoder_config
 
     if custom_config is not None:
         pretrained_config.update(custom_config)
@@ -550,6 +569,7 @@ def load_pretrained_model(
         tag=model_metadata["net"].get("tag", None),
         image_encoder=image_encoder,
         text_encoder=text_encoder,
+        text_decoder=text_decoder,
         embed_dim=model_metadata["net"].get("embed_dim", None),
         tokenizer=model_metadata["net"].get("tokenizer", None),
         inference=inference,

{birder_clip-0.0.2.dev7 → birder_clip-0.0.2.dev8}/birder_clip/common/lib.py

@@ -3,14 +3,17 @@ from typing import Any
 from typing import Optional
 
 from birder.data.transforms.classification import RGBType
+from birder.version import __version__ as birder_version
 
 from birder_clip.conf import settings
 from birder_clip.model_registry import registry
 from birder_clip.net.base import BaseNet
 from birder_clip.net.base import SignatureType
-from birder_clip.version import __version__
+from birder_clip.version import __version__ as birder_clip_version
 
-MODEL_CONFIG_RESERVED_KEYS = frozenset({"image", "text", "tokenizer", "embed_dim", "embed-dim"})
+MODEL_CONFIG_RESERVED_KEYS = frozenset(
+    {"image", "text", "decoder", "tokenizer", "embed_dim", "embed-dim", "keep_ratio"}
+)
 
 
 def get_size_from_signature(signature: SignatureType) -> tuple[int, int]:
@@ -37,6 +40,7 @@ def get_image_text_network_name(
     tag: Optional[str] = None,
     image_encoder: Optional[str] = None,
     text_encoder: Optional[str] = None,
+    text_decoder: Optional[str] = None,
     embed_dim: Optional[int] = None,
     tokenizer: Optional[str] = None,
 ) -> str:
@@ -45,6 +49,8 @@ def get_image_text_network_name(
         parts.append(image_encoder)
     if text_encoder is not None and text_encoder != "transformer_encoder":
         parts.append(text_encoder)
+    if text_decoder is not None and text_decoder != "conditioned_decoder":
+        parts.append(text_decoder)
 
     if registry.exists(network) is True:
         default_tokenizer = registry.get_default_tokenizer(network)
@@ -71,10 +77,12 @@ def get_image_text_model_config(
     *,
     image_encoder: Optional[str] = None,
     text_encoder: Optional[str] = None,
+    text_decoder: Optional[str] = None,
     embed_dim: Optional[int] = None,
     tokenizer: Optional[str] = None,
     image_encoder_config: Optional[dict[str, Any]] = None,
     text_encoder_config: Optional[dict[str, Any]] = None,
+    text_decoder_config: Optional[dict[str, Any]] = None,
     input_channels: Optional[int] = None,
     image_size: Optional[tuple[int, int]] = None,
     context_length: Optional[int] = None,
@@ -86,7 +94,7 @@ def get_image_text_model_config(
 
     if config is not None:
         for key, value in config.items():
-            if key in {"image", "text"} and isinstance(value, dict):
+            if key in {"image", "text", "decoder"} and isinstance(value, dict):
                 model_config[key] = {**model_config.get(key, {}), **value}
             else:
                 model_config[key] = value
@@ -111,7 +119,7 @@ def get_image_text_model_config(
 
         model_config["image"] = image_config
 
-    if text_encoder is not None or text_encoder_config is not None or context_length is not None:
+    if text_encoder is not None or text_encoder_config is not None or "text" in model_config:
         text_config = model_config.get("text", {}).copy()
         if text_encoder is not None:
             # String encoder metadata replaces only the encoder name.
@@ -124,6 +132,17 @@ def get_image_text_model_config(
 
         model_config["text"] = text_config
 
+    if text_decoder is not None or text_decoder_config is not None or "decoder" in model_config:
+        decoder_config = model_config.get("decoder", {}).copy()
+        if text_decoder is not None:
+            decoder_config["network"] = text_decoder
+        if text_decoder_config is not None:
+            decoder_config["config"] = {**decoder_config.get("config", {}), **text_decoder_config}
+        if context_length is not None:
+            decoder_config["context_length"] = context_length
+
+        model_config["decoder"] = decoder_config
+
     if embed_dim is not None:
         model_config["embed_dim"] = embed_dim
     if tokenizer is not None:
@@ -143,7 +162,8 @@ def get_image_text_network_config(net: BaseNet, signature: SignatureType, rgb_st
         model_config = net.config
 
     return {
-        "birder_clip_version": __version__,
+        "birder_clip_version": birder_clip_version,
+        "birder_version": birder_version,
         "name": model_name,
         "registered_name": registered_name,
         "task": net.task,

{birder_clip-0.0.2.dev7 → birder_clip-0.0.2.dev8}/birder_clip/common/training_cli.py

@@ -27,6 +27,7 @@ def add_model_args(parser: argparse.ArgumentParser) -> None:
         help="pretrained Birder image model weights path to load into the image encoder",
     )
     parser.add_argument("--text-encoder", type=str, help="the text encoder to use")
+    parser.add_argument("--text-decoder", type=str, help="the text decoder to use")
     parser.add_argument("--embed-dim", type=int, metavar="N", help="shared image-text embedding dimension")
     parser.add_argument("--tokenizer", type=str, help="the tokenizer to use")
     parser.add_argument(
@@ -44,11 +45,21 @@ def add_model_args(parser: argparse.ArgumentParser) -> None:
         action=cli.FlexibleDictAction,
         help="override the text encoder configuration, accepts key-value pairs or JSON",
     )
+    parser.add_argument(
+        "--text-decoder-config",
+        action=cli.FlexibleDictAction,
+        help="override the text decoder configuration, accepts key-value pairs or JSON",
+    )
+    parser.add_argument(
+        "--openvision-v2-keep-ratio", type=float, help="OpenVision v2 image token keep ratio for caption decoding"
+    )
 
 
 def add_loss_args(parser: argparse.ArgumentParser) -> None:
     group = parser.add_argument_group("Loss parameters")
-    group.add_argument("--loss", type=str, choices=["clip", "coca"], default="clip", help="loss function to use")
+    group.add_argument(
+        "--loss", type=str, choices=["clip", "coca", "caption"], default="clip", help="loss function to use"
+    )
     group.add_argument(
         "--coca-caption-loss-weight", type=float, default=1.0, help="weight assigned to CoCa caption loss"
     )
@@ -662,13 +673,16 @@ def common_args_validation(args: argparse.Namespace) -> None:
         raise cli.ValidationError("--grad-accum-steps must be >= 1")
     if args.grad_accum_cache_negatives is True and args.grad_accum_steps == 1:
         raise cli.ValidationError("--grad-accum-cache-negatives requires --grad-accum-steps greater than 1")
-    if args.grad_accum_cache_negatives is True and args.loss == "coca":
+    if args.grad_accum_cache_negatives is True and args.loss != "clip":
         raise cli.ValidationError("--grad-accum-cache-negatives is only supported with --loss clip")
 
     if args.coca_caption_loss_weight < 0.0:
         raise cli.ValidationError("--coca-caption-loss-weight must be non-negative")
     if args.coca_contrastive_loss_weight < 0.0:
         raise cli.ValidationError("--coca-contrastive-loss-weight must be non-negative")
+    if args.openvision_v2_keep_ratio is not None:
+        if args.openvision_v2_keep_ratio <= 0.0 or args.openvision_v2_keep_ratio > 1.0:
+            raise cli.ValidationError("--openvision-v2-keep-ratio must be in range of (0, 1]")
 
     # EMA
     if args.model_ema_steps < 1:

{birder_clip-0.0.2.dev7 → birder_clip-0.0.2.dev8}/birder_clip/common/training_utils.py

@@ -1,4 +1,5 @@
 import argparse
+import json
 import logging
 from pathlib import Path
 from typing import Any
@@ -11,6 +12,7 @@ from birder.common import training_utils as birder_training_utils
 
 from birder_clip.common import fs_ops
 from birder_clip.conf import settings
+from birder_clip.version import __version__ as birder_clip_version
 
 
 def setup_file_logging(log_file_path: str | Path) -> logging.Handler:
@@ -28,6 +30,18 @@ def setup_file_logging(log_file_path: str | Path) -> logging.Handler:
     return file_handler
 
 
+def make_training_args_payload(args: argparse.Namespace) -> dict[str, Any]:
+    return {
+        "birder_clip_version": birder_clip_version,
+        **birder_training_utils.make_training_args_payload(args),
+    }
+
+
+def write_training_args_json(path: Path, args: argparse.Namespace) -> None:
+    with open(path.joinpath("training_args.json"), "w", encoding="utf-8") as handle:
+        json.dump(make_training_args_payload(args), handle, indent=2)
+
+
 def save_training_checkpoint(
     args: argparse.Namespace,
     network_name: str,

{birder_clip-0.0.2.dev7 → birder_clip-0.0.2.dev8}/birder_clip/loss/__init__.py

@@ -1,7 +1,9 @@
+from birder_clip.loss.caption import CaptionLoss
 from birder_clip.loss.coca import CoCaLoss
 from birder_clip.loss.contrastive import CLIPLoss
 
 __all__ = [
+    "CaptionLoss",
     "CoCaLoss",
     "CLIPLoss",
 ]

birder_clip-0.0.2.dev8/birder_clip/loss/caption.py

@@ -0,0 +1,72 @@
+from collections.abc import Sequence
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+
+class CaptionLoss(torch.nn.Module):
+    """
+    Autoregressive captioning cross entropy over decoder logits
+
+    The loss consumes unshifted tokenized captions and supports two decoder output conventions:
+    - logits length equals text length: the final logit is ignored.
+    - logits length equals text length - 1: logits are used as-is.
+    """
+
+    def __init__(
+        self,
+        caption_loss_weight: float = 1.0,
+        pad_token_id: int = 0,
+        ignore_token_ids: Optional[Sequence[int]] = None,
+        label_smoothing: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.caption_loss_weight = caption_loss_weight
+        self.ignore_token_ids = tuple(ignore_token_ids) if ignore_token_ids is not None else (pad_token_id,)
+        self.label_smoothing = label_smoothing
+
+    def _align_logits_and_targets(self, logits: torch.Tensor, texts: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        if logits.size(1) == texts.size(1):
+            return logits[:, :-1], texts[:, 1:]
+        if logits.size(1) == texts.size(1) - 1:
+            return logits, texts[:, 1:]
+
+        raise ValueError(
+            "Expected logits sequence length to equal text sequence length or text sequence length - 1, "
+            f"got logits={logits.size(1)}, texts={texts.size(1)}"
+        )
+
+    def _align_target_mask(self, target_mask: torch.Tensor, texts: torch.Tensor) -> torch.Tensor:
+        if target_mask.shape == texts.shape:
+            target_mask = target_mask[:, 1:]
+
+        return target_mask
+
+    def forward(
+        self, logits: torch.Tensor, texts: torch.Tensor, target_mask: Optional[torch.Tensor] = None
+    ) -> dict[str, torch.Tensor]:
+        if self.caption_loss_weight == 0.0:
+            return {"caption_loss": logits.new_zeros(())}
+
+        logits, targets = self._align_logits_and_targets(logits, texts)
+        token_loss = F.cross_entropy(
+            logits.permute(0, 2, 1),
+            targets,
+            reduction="none",
+            label_smoothing=self.label_smoothing,
+        )
+
+        loss_mask = torch.ones_like(targets, dtype=torch.bool)
+        for token_id in self.ignore_token_ids:
+            loss_mask = loss_mask & (targets != token_id)
+
+        if target_mask is not None:
+            target_mask = self._align_target_mask(target_mask, texts)
+            loss_mask = loss_mask & target_mask.to(dtype=torch.bool)
+
+        loss_mask_float = loss_mask.to(dtype=token_loss.dtype)
+        caption_loss = (token_loss * loss_mask_float).sum() / loss_mask_float.sum().clamp_min(1)
+        caption_loss = caption_loss * self.caption_loss_weight
+
+        return {"caption_loss": caption_loss}

{birder_clip-0.0.2.dev7 → birder_clip-0.0.2.dev8}/birder_clip/loss/coca.py

@@ -26,9 +26,7 @@ class CoCaLoss(torch.nn.Module):
     captioning cross entropy over decoder logits.
     """
 
-    def __init__(
-        self, *, caption_loss_weight: float = 1.0, clip_loss_weight: float = 1.0, pad_token_id: int = 0
-    ) -> None:
+    def __init__(self, caption_loss_weight: float = 1.0, clip_loss_weight: float = 1.0, pad_token_id: int = 0) -> None:
         super().__init__()
         self.caption_loss_weight = caption_loss_weight
         self.clip_loss_weight = clip_loss_weight

{birder_clip-0.0.2.dev7 → birder_clip-0.0.2.dev8}/birder_clip/model_registry/manifest.py

@@ -11,8 +11,8 @@ FormatInfoType = TypedDict(
     {"file_size": float, "sha256": str},
 )
 
-EncoderInfoType = TypedDict(
-    "EncoderInfoType",
+ImageEncoderInfoType = TypedDict(
+    "ImageEncoderInfoType",
     {
         "network": str,
         "config": NotRequired[dict[str, Any]],
@@ -21,16 +21,33 @@ EncoderInfoType = TypedDict(
         "size": NotRequired[tuple[int, int]],
     },
 )
+TextEncoderInfoType = TypedDict(
+    "TextEncoderInfoType",
+    {
+        "network": str,
+        "config": NotRequired[dict[str, Any]],
+    },
+)
+TextDecoderInfoType = TypedDict(
+    "TextDecoderInfoType",
+    {
+        "network": str,
+        "config": NotRequired[dict[str, Any]],
+    },
+)
 
-EncoderMetadataType: TypeAlias = str | EncoderInfoType
+ImageEncoderMetadataType: TypeAlias = str | ImageEncoderInfoType
+TextEncoderMetadataType: TypeAlias = str | TextEncoderInfoType
+TextDecoderMetadataType: TypeAlias = str | TextDecoderInfoType
 
 NetworkInfoType = TypedDict(
     "NetworkInfoType",
     {
         "network": str,
         "tag": NotRequired[str],
-        "image_encoder": NotRequired[EncoderMetadataType],
-        "text_encoder": NotRequired[EncoderMetadataType],
+        "image_encoder": NotRequired[ImageEncoderMetadataType],
+        "text_encoder": NotRequired[TextEncoderMetadataType],
+        "text_decoder": NotRequired[TextDecoderMetadataType],
         "embed_dim": NotRequired[int],
         "tokenizer": NotRequired[str],
     },

{birder_clip-0.0.2.dev7 → birder_clip-0.0.2.dev8}/birder_clip/model_registry/model_registry.py

@@ -170,6 +170,11 @@ class ModelRegistry:
             if context_length is not None:
                 return context_length  # type: ignore[no-any-return]
 
+            decoder_config = config.get("decoder", {})
+            context_length = decoder_config.get("context_length")
+            if context_length is not None:
+                return context_length  # type: ignore[no-any-return]
+
             if tokenizer is None:
                 tokenizer = config.get("tokenizer")
 

{birder_clip-0.0.2.dev7 → birder_clip-0.0.2.dev8}/birder_clip/net/__init__.py

@@ -1,7 +1,9 @@
 from birder_clip.net.clip import CLIP
 from birder_clip.net.coca import CoCa
+from birder_clip.net.openvision_v2 import OpenVision_v2
 
 __all__ = [
     "CLIP",
     "CoCa",
+    "OpenVision_v2",
 ]