Add source code

src/open_clip/.ipynb_checkpoints/tokenizer-checkpoint.py

@@ -0,0 +1,621 @@
+""" CLIP tokenizer
+
+Copied from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+import gzip
+import html
+import os
+import random
+import string
+from functools import lru_cache, partial
+from typing import Callable, List, Optional, Union, Dict
+import warnings
+
+import ftfy
+import numpy as np
+import regex as re
+import torch
+
+# https://stackoverflow.com/q/62691279
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+_nltk_init = False
+
+DEFAULT_CONTEXT_LENGTH = 77  # default context length for OpenAI CLIP
+
+
+@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 significant 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 = " ".join(text.split())
+    text = text.strip()
+    return text
+
+
+def _clean_canonicalize(x):
+    # basic, remove whitespace, remove punctuation, lower case
+    return canonicalize_text(basic_clean(x))
+
+
+def _clean_lower(x):
+    # basic, remove whitespace, lower case
+    return whitespace_clean(basic_clean(x)).lower()
+
+
+def _clean_whitespace(x):
+    # basic, remove whitespace
+    return whitespace_clean(basic_clean(x))
+
+
+def get_clean_fn(type: str):
+    if type == 'canonicalize':
+        return _clean_canonicalize
+    elif type == 'lower':
+        return _clean_lower
+    elif type == 'whitespace':
+        return _clean_whitespace
+    else:
+        assert False, f"Invalid clean function ({type})."
+
+
+def canonicalize_text(
+    text,
+    *,
+    keep_punctuation_exact_string=None,
+    trans_punctuation: dict = str.maketrans("", "", string.punctuation),
+):
+    """Returns canonicalized `text` (lowercase and punctuation removed).
+
+    From: https://github.com/google-research/big_vision/blob/53f18caf27a9419231bbf08d3388b07671616d3d/big_vision/evaluators/proj/image_text/prompt_engineering.py#L94
+
+    Args:
+      text: string to be canonicalized.
+      keep_punctuation_exact_string: If provided, then this exact string kept.
+        For example providing '{}' will keep any occurrences of '{}' (but will
+        still remove '{' and '}' that appear separately).
+    """
+    text = text.replace("_", " ")
+    if keep_punctuation_exact_string:
+        text = keep_punctuation_exact_string.join(
+            part.translate(trans_punctuation)
+            for part in text.split(keep_punctuation_exact_string)
+        )
+    else:
+        text = text.translate(trans_punctuation)
+    text = text.lower()
+    text = " ".join(text.split())
+    return text.strip()
+
+
+class SimpleTokenizer(object):
+    def __init__(
+            self,
+            bpe_path: str = default_bpe(),
+            additional_special_tokens: Optional[List[str]] = None,
+            context_length: Optional[int] = DEFAULT_CONTEXT_LENGTH,
+            clean: str = 'lower',
+            reduction_mask: str = ''
+    ):
+        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))
+        special_tokens = ['<start_of_text>', '<end_of_text>']
+        if additional_special_tokens:
+            special_tokens += additional_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]
+        self.sot_token_id = self.all_special_ids[0]
+        self.eot_token_id = self.all_special_ids[1]
+        self.context_length = context_length
+        self.clean_fn = get_clean_fn(clean)
+        self.reduction_fn = get_reduction_mask_fn(reduction_mask) if reduction_mask else None
+
+    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 Exception:
+                    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 = self.clean_fn(text)
+        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
+
+    def __call__(self, texts: Union[str, List[str]], context_length: Optional[int] = None) -> 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]
+
+        context_length = context_length or self.context_length
+        assert context_length, 'Please set a valid context length'
+
+        if self.reduction_fn is not None:
+            # use reduction strategy for tokenize if set, otherwise default to truncation below
+            return self.reduction_fn(
+                texts,
+                context_length=context_length,
+                sot_token_id=self.sot_token_id,
+                eot_token_id=self.eot_token_id,
+                encode_fn=self.encode,
+            )
+
+        all_tokens = [[self.sot_token_id] + self.encode(text) + [self.eot_token_id] 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:
+                tokens = tokens[:context_length]  # Truncate
+                tokens[-1] = self.eot_token_id
+            result[i, :len(tokens)] = torch.tensor(tokens)
+
+        return result
+
+
+_tokenizer = SimpleTokenizer()
+
+
+def decode(output_ids: torch.Tensor):
+    output_ids = output_ids.cpu().numpy()
+    return _tokenizer.decode(output_ids)
+
+
+def tokenize(texts: Union[str, List[str]], context_length: int = DEFAULT_CONTEXT_LENGTH) -> torch.LongTensor:
+    return _tokenizer(texts, context_length=context_length)
+
+
+def random_mask_tokenize(
+        texts: Union[str, List[str]],
+        context_length: int,
+        sot_token_id: int,
+        eot_token_id: int,
+        encode_fn: Callable,
+        shuffle: bool = False,
+):
+    all_tokens = [encode_fn(text) for text in texts]
+    result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+    for i, tokens in enumerate(all_tokens):
+        tokens = torch.tensor(tokens)
+        num_tokens = len(tokens)
+        if num_tokens > context_length - 2:  # 2 for sot and eot token
+            num_keep = context_length - 2
+            indices = torch.randperm(len(tokens))
+            indices = indices[:num_keep]
+            if not shuffle:
+                indices = indices.msort()
+            tokens = tokens[indices]
+            num_tokens = num_keep
+        result[i, 0] = sot_token_id
+        result[i, 1:num_tokens + 1] = tokens
+        result[i, num_tokens + 1] = eot_token_id
+
+    return result
+
+
+def simple_mask_tokenize(
+        texts: Union[str, List[str]],
+        context_length: int,
+        sot_token_id: int,
+        eot_token_id: int,
+        encode_fn: Callable,
+):
+    all_tokens = [encode_fn(text) for text in texts]
+    result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+    for i, tokens in enumerate(all_tokens):
+        num_tokens = len(tokens)
+        if num_tokens > context_length - 2:  # 2 for sot and eot token
+            num_keep = context_length - 2
+            start_index = random.randint(0, num_tokens - num_keep)  # high is incl
+            tokens = tokens[start_index: start_index + num_keep]
+        tokens = [sot_token_id] + tokens + [eot_token_id]
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result
+
+
+def syntax_mask_tokenize(
+        texts: Union[str, List[str]],
+        context_length: int,
+        sot_token_id: int,
+        eot_token_id: int,
+        encode_fn: Callable,
+) -> torch.LongTensor:
+    """ Returns the tokenized representation of given input string(s).
+    Apply syntax masking before tokenize.
+    """
+    import nltk
+    global _nltk_init
+    if not _nltk_init:
+        # run them for the first time
+        nltk.download('punkt')
+        nltk.download('averaged_perceptron_tagger')
+        _nltk_init = True
+
+    def get_order(x):
+        if x.startswith('NN'):
+            return 1
+        elif x.startswith('JJ'):
+            return 2
+        elif x.startswith('VB'):
+            return 3
+        else:
+            return 4
+
+    # syntax masking
+    new_texts = []
+    for text in texts:
+        list_tokens = nltk.tokenize.word_tokenize(text)
+        pos_tags = nltk.pos_tag(list_tokens)
+        #  sample the words by get_order method
+        order_list = [get_order(tag) for _, tag in pos_tags]
+        sorted_ids = np.argsort(np.array(order_list))
+        sampled_ids = sorted(sorted_ids[:context_length - 2]) # need 2 slots for sot and eot tokens
+        sampled_tokens = np.take(np.array(list_tokens), sampled_ids, axis=0)  # sample the tokens
+
+        new_text = ''
+        for token in sampled_tokens:
+            new_text = new_text + str(token) + ' '
+        new_text = new_text.strip()
+        new_texts.append(new_text)
+    texts = new_texts
+
+    all_tokens = [[sot_token_id] + encode_fn(text) + [eot_token_id] for text in texts]
+    result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+    for i, tokens in enumerate(all_tokens):
+        # still need first truncate because some words produces two tokens
+        if len(tokens) > context_length:
+            tokens = tokens[:context_length]  # Truncate
+            tokens[-1] = eot_token_id
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result
+
+
+def get_reduction_mask_fn(type: str):
+    """ Choose strategy for dropping (masking) tokens to achieve target context length"""
+    assert type in ('simple', 'random', 'shuffle', 'syntax')
+    if type == 'simple':
+        return simple_mask_tokenize  # randomly select block [start:end]
+    elif type == 'random':
+        return random_mask_tokenize  # randomly drop tokens (keep order)
+    elif type == 'shuffle':
+        return partial(random_mask_tokenize, shuffle=True)  # randomly drop tokens (shuffle order)
+    elif type == 'syntax':
+        return syntax_mask_tokenize  # randomly drop prioritized by syntax
+    else:
+        assert False, F'Unknown type {type}.'
+
+
+class HFTokenizer:
+    """HuggingFace tokenizer wrapper with support for custom tokenization modes"""
+
+    def __init__(
+            self,
+            tokenizer_name: str,
+            context_length: Optional[int] = DEFAULT_CONTEXT_LENGTH,
+            clean: str = 'whitespace',
+            strip_sep_token: bool = False,
+            language: Optional[str] = None,
+            cache_dir: Optional[str] = None,
+            tokenizer_mode: Optional[str] = None,  # None, 'clips'
+            **kwargs
+    ):
+        self.tokenizer_mode = tokenizer_mode or ''
+        self.context_length = context_length
+        self.clean_fn = get_clean_fn(clean)
+        self.strip_sep_token = strip_sep_token
+
+        # NOTE: Left as example of loading custom tokenizer from file for experimentation
+        # if self.tokenizer_mode == 'bert_clips':
+        #     self.special_tokens = {
+        #         "bos_token": 1,
+        #         "eos_token": 2,
+        #         "cls_token": 101,
+        #         "pad_token": 0
+        #     }
+        #
+        #     # For BERT CLIPS mode with vocab file
+        #     from tokenizers import BertWordPieceTokenizer
+        #     if tokenizer_name.startswith('hf-hub:'):
+        #         from huggingface_hub import hf_hub_download
+        #         # Format: hf-hub:repo_id/filename
+        #         repo_url = tokenizer_name[7:]
+        #         parts = repo_url.split('/')
+        #         filename = parts[-1]
+        #         repo_id = '/'.join(parts[:-1])
+        #         vocab_file = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=cache_dir)
+        #         self.tokenizer = BertWordPieceTokenizer(lowercase=True)
+        #         self.tokenizer = self.tokenizer.from_file(vocab_file)
+        #     else:
+        #         # Assume tokenizer_name is a local path to a vocab file
+        #         self.tokenizer = BertWordPieceTokenizer(lowercase=True)
+        #         self.tokenizer = self.tokenizer.from_file(tokenizer_name)
+
+        # Standard HuggingFace tokenizer initialization
+        from transformers import AutoTokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_name,
+            cache_dir=cache_dir,
+            **kwargs
+        )
+
+        # Set language function if available
+        set_lang_fn = getattr(self.tokenizer, 'set_src_lang_special_tokens', None)
+        if callable(set_lang_fn):
+            self.set_lang_fn = set_lang_fn
+        if language is not None:
+            self.set_language(language)
+
+    def save_pretrained(self, dest):
+        self.tokenizer.save_pretrained(dest)
+
+    def __call__(self, texts: Union[str, List[str]], context_length: Optional[int] = None) -> torch.Tensor:
+        # same cleaning as for default tokenizer, except lowercasing
+        # adding lower (for case-sensitive tokenizers) will make it more robust but less sensitive to nuance
+        if isinstance(texts, str):
+            texts = [texts]
+
+        context_length = context_length or self.context_length
+        assert context_length, 'Please set a valid context length in class init or call.'
+
+        texts = [self.clean_fn(text) for text in texts]
+
+        # Handle different tokenization modes
+        if self.tokenizer_mode == 'clips':
+            return self._clips_tokenize(texts, context_length)
+        else:
+            # Standard tokenization
+            input_ids = self.tokenizer.batch_encode_plus(
+                texts,
+                return_tensors='pt',
+                max_length=context_length,
+                padding='max_length',
+                truncation=True,
+            ).input_ids
+
+            if self.strip_sep_token:
+                input_ids = torch.where(
+                    input_ids == self.tokenizer.sep_token_id,
+                    torch.zeros_like(input_ids),
+                    input_ids,
+                )
+
+            return input_ids
+
+    def set_language(self, src_lang):
+        if hasattr(self, 'set_lang_fn'):
+            self.set_lang_fn(src_lang)
+        else:
+            warnings.warn('Cannot set language for the tokenizer.')
+
+    def _clips_tokenize(self, texts: List[str], context_length: int) -> torch.Tensor:
+        """Use standard HF tokenizer but apply custom post-processing"""
+        # Use standard tokenizer without special tokens - we'll add our own
+        encoded_outputs = self.tokenizer.batch_encode_plus(
+            texts,
+            add_special_tokens=False,
+            padding=False,
+            truncation=False,
+            return_tensors=None
+        )
+
+        encoded = []
+        for tokens in encoded_outputs["input_ids"]:
+            tokens = tokens[:context_length - 3]  # Leave room for special tokens
+            tokens = [self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id]
+            encoded.append(tokens)
+
+        # Create result tensor and handle padding + class token
+        result = torch.zeros(len(encoded), context_length, dtype=torch.long)
+        for i, tokens in enumerate(encoded):
+            padded_tokens = self._pad_and_add_class_token(
+                tokens,
+                max_length=context_length,
+                pad_token_id=self.tokenizer.pad_token_id,
+                cls_token_id=self.tokenizer.cls_token_id,
+            )
+            result[i, :len(padded_tokens)] = torch.tensor(padded_tokens)
+
+        return result
+
+    def _pad_and_add_class_token(
+            self,
+            tokens: List[int],
+            max_length: int,
+            pad_token_id: int = 0,
+            cls_token_id: int = 101,
+    ) -> List[int]:
+        """ Add padding with class token at the end """
+        if len(tokens) > max_length - 1:
+            tokens = tokens[:max_length - 1]
+
+        # Add padding to reach max_length-1
+        if len(tokens) < max_length - 1:
+            tokens = tokens + [pad_token_id] * (max_length - 1 - len(tokens))
+
+        # Add class token at the end
+        tokens = tokens + [cls_token_id]
+        return tokens
+
+
+class SigLipTokenizer:
+    """HuggingFace tokenizer wrapper for SigLIP T5 compatible sentencepiece vocabs
+
+    NOTE: this is not needed in normal library use, but is used to import new sentencepiece tokenizers
+    into OpenCLIP. Leaving code here in case future models use new tokenizers.
+    """
+    VOCAB_FILES = {
+        # english, vocab_size=32_000
+        "c4-en": "http://storage.googleapis.com/t5-data/vocabs/cc_en.32000/sentencepiece.model",
+        # used in multilingual models (mT5, PaLI), vocab_size=250_000
+        "mc4": "http://storage.googleapis.com/t5-data/vocabs/mc4.250000.100extra/sentencepiece.model",
+        # used in SigLIP2 models, vocab_size=256000
+        "gemma": "http://storage.googleapis.com/big_vision/gemma_tokenizer.model",
+    }
+
+    def __init__(
+            self,
+            tokenizer_name: str,
+            context_length: Optional[int] = 64,
+    ):
+        if 'gemma' in tokenizer_name:
+            from transformers import GemmaTokenizerFast
+            tokenizer_cls = partial(
+                GemmaTokenizerFast, padding_side='right', add_bos_token=False, add_eos_token=True)
+        else:
+            from transformers import T5TokenizerFast
+            tokenizer_cls = partial(T5TokenizerFast, extra_ids=0)
+
+        if tokenizer_name in self.VOCAB_FILES:
+            # FIXME temporary hack?
+            import tempfile
+            import fsspec
+            vocab_file = self.VOCAB_FILES[tokenizer_name]
+            with tempfile.NamedTemporaryFile('wb') as dst:
+                with fsspec.open(vocab_file, 'rb') as src:
+                    dst.write(src.read())
+                self.tokenizer = tokenizer_cls(dst.name, legacy=False)
+        else:
+            self.tokenizer = tokenizer_cls(tokenizer_name, legacy=False)
+
+        self.tokenizer.pad_token_id = 0 if 'gemma' in tokenizer_name else 1
+        self.tokenizer.eos_token_id = 1
+        self.context_length = context_length
+
+    def save_pretrained(self, dest):
+        self.tokenizer.save_pretrained(dest)
+
+    def __call__(self, texts: Union[str, List[str]], context_length: Optional[int] = None) -> torch.Tensor:
+        # same cleaning as for default tokenizer, except lowercasing
+        # adding lower (for case-sensitive tokenizers) will make it more robust but less sensitive to nuance
+        if isinstance(texts, str):
+            texts = [texts]
+
+        context_length = context_length or self.context_length
+        assert context_length, 'Please set a valid context length in class init or call.'
+
+        texts = [canonicalize_text(basic_clean(text)) for text in texts]
+        output = self.tokenizer(
+            texts,
+            return_tensors='pt',
+            max_length=context_length,
+            padding='max_length',
+            truncation=True,
+        )
+        return output.input_ids

src/open_clip/__init__.py

@@ -0,0 +1,2 @@
+from .factory import create_model, load_checkpoint, get_tokenizer, get_input_dtype
+from .tokenizer import SimpleTokenizer

src/open_clip/biosignals_coca_model.py

@@ -0,0 +1,1807 @@
+"""
+Biosignals-Text CoCa Model
+
+Adapted from the original CoCa model to work with biosignals (time series) data
+instead of images. This model is designed for biosignals-text contrastive learning.
+"""
+
+from typing import Dict, List, Optional, Union, Tuple
+import torch
+from torch import nn
+from torch.nn import functional as F
+import numpy as np
+import math
+from dataclasses import dataclass, field
+
+from .transformer import (
+    LayerNormFp32,
+    LayerNorm,
+    QuickGELU,
+    MultimodalTransformer,
+    ConcatMultimodalTransformer,
+)
+from .model import CLIPTextCfg, _build_text_tower
+from .coca_model import MultimodalCfg, _build_text_decoder_tower, _token_to_tensor
+
+try:
+    from transformers.generation.beam_search import BeamSearchScorer
+    from transformers.generation.logits_process import (
+        LogitsProcessorList,
+        TopPLogitsWarper,
+        TopKLogitsWarper,
+        RepetitionPenaltyLogitsProcessor,
+        MinLengthLogitsProcessor,
+    )
+    from transformers.generation.stopping_criteria import (
+        MaxLengthCriteria,
+        EosTokenCriteria,
+        StoppingCriteriaList,
+    )
+
+    GENERATION_TYPES = {
+        "top_k": TopKLogitsWarper,
+        "top_p": TopPLogitsWarper,
+        "beam_search": "beam_search"
+    }
+    _has_transformers = True
+except ImportError as e:
+    GENERATION_TYPES = {
+        "top_k": None,
+        "top_p": None,
+        "beam_search": "beam_search"
+    }
+    _has_transformers = False
+
+
+# ============================================================================
+# Pure Transformer Architecture Components (from PureTransformerMAE)
+# ============================================================================
+
+class RotaryEmbedding(nn.Module):
+    """Rotary Position Embedding (RoPE)"""
+    def __init__(self, dim: int, theta: float = 10000.0, learned_freq: bool = False):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.learned_freq = learned_freq
+        
+        if learned_freq:
+            # Learnable frequencies for channel attention
+            self.freqs = nn.Parameter(torch.randn(dim // 2) * 0.02)
+        else:
+            # Fixed frequencies for temporal attention
+            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+            self.register_buffer('freqs', freqs)
+    
+    def rotate_queries_or_keys(self, x: torch.Tensor, position_ids: Optional[torch.Tensor] = None):
+        """
+        Apply rotary embeddings to queries or keys
+        
+        Args:
+            x: (batch_size, num_heads, seq_len, head_dim)
+            position_ids: (seq_len,) or (batch_size, seq_len) - position indices
+        Returns:
+            Rotated tensor of same shape
+        """
+        batch_size, num_heads, seq_len, head_dim = x.shape
+        assert head_dim == self.dim, f"head_dim {head_dim} != self.dim {self.dim}"
+        
+        # Generate position indices if not provided
+        if position_ids is None:
+            position_ids = torch.arange(seq_len, device=x.device, dtype=torch.float)
+        elif position_ids.ndim == 2:
+            # If 2D, take the first batch (assuming all batches have same pattern)
+            position_ids = position_ids[0].float()
+        else:
+            position_ids = position_ids.float()
+        
+        # Compute angles: position_ids * freqs
+        # position_ids: (seq_len,), freqs: (dim // 2,)
+        # angles: (seq_len, dim // 2)
+        angles = torch.einsum('s,d->sd', position_ids, self.freqs)
+        
+        # Duplicate for cos and sin
+        # cos/sin: (seq_len, dim)
+        cos = torch.cos(angles).repeat_interleave(2, dim=-1)
+        sin = torch.sin(angles).repeat_interleave(2, dim=-1)
+        
+        # Reshape for broadcasting: (1, 1, seq_len, dim)
+        cos = cos.unsqueeze(0).unsqueeze(0)
+        sin = sin.unsqueeze(0).unsqueeze(0)
+        
+        # Apply rotation
+        # Split x into even and odd dimensions
+        x1 = x[..., 0::2]  # Even dimensions
+        x2 = x[..., 1::2]  # Odd dimensions
+        
+        # Apply rotation: [x1, x2] @ [[cos, -sin], [sin, cos]]
+        x_rotated = torch.empty_like(x)
+        x_rotated[..., 0::2] = x1 * cos[..., 0::2] - x2 * sin[..., 0::2]
+        x_rotated[..., 1::2] = x1 * sin[..., 1::2] + x2 * cos[..., 1::2]
+        
+        return x_rotated
+
+
+class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization"""
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU activation function: SiLU(x * W1) * (x * W2)"""
+    def __init__(self, dim_in: int, dim_out: int, bias: bool = False):
+        super().__init__()
+        self.w1 = nn.Linear(dim_in, dim_out, bias=bias)
+        self.w2 = nn.Linear(dim_in, dim_out, bias=bias)
+        
+    def forward(self, x):
+        return F.silu(self.w1(x)) * self.w2(x)
+
+
+class MLP(nn.Module):
+    """MLP with configurable activation and normalization"""
+    def __init__(self, 
+                 dim: int, 
+                 hidden_dim: int, 
+                 dropout: float = 0.0,
+                 activation: str = "swiglu",  # "swiglu", "gelu", "relu"
+                 bias: bool = False):
+        super().__init__()
+        self.activation = activation
+        
+        if activation == "swiglu":
+            # SwiGLU requires different structure: two parallel linear layers
+            self.gate_proj = SwiGLU(dim, hidden_dim, bias=bias)
+            self.down_proj = nn.Linear(hidden_dim, dim, bias=bias)
+        else:
+            # Standard MLP structure
+            self.up_proj = nn.Linear(dim, hidden_dim, bias=bias)
+            self.down_proj = nn.Linear(hidden_dim, dim, bias=bias)
+            
+            if activation == "gelu":
+                self.act_fn = nn.GELU()
+            elif activation == "relu":
+                self.act_fn = nn.ReLU()
+            else:
+                raise ValueError(f"Unknown activation: {activation}")
+        
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, x):
+        if self.activation == "swiglu":
+            x = self.gate_proj(x)
+            x = self.dropout(x)
+            x = self.down_proj(x)
+        else:
+            x = self.up_proj(x)
+            x = self.act_fn(x)
+            x = self.dropout(x)
+            x = self.down_proj(x)
+            
+        return self.dropout(x)
+
+
+class ChannelPatching(nn.Module):
+    """Patching layer that operates independently on each channel"""
+    def __init__(self, 
+                 patch_size: int = 32,
+                 conv_embed_dim: int = 256,
+                 num_channels: int = 21):
+        super().__init__()
+        self.patch_size = patch_size
+        self.conv_embed_dim = conv_embed_dim
+        self.num_channels = num_channels
+        
+        # Single conv layer applied to all channels (kernel_size=patch_size, stride=patch_size)
+        self.conv_patching = nn.Conv1d(
+            in_channels=1, 
+            out_channels=conv_embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            padding=0  # No padding for clean non-overlapping patches
+        )
+        
+    def forward(self, x):
+        """
+        Args:
+            x: (batch_size, num_channels, signal_length) - multi-channel signal
+        Returns:
+            (batch_size, num_channels, num_patches, conv_embed_dim) - patched representations
+        """
+        batch_size, num_channels, seq_len = x.shape
+        
+        # Reshape to process all channels independently: (batch_size * num_channels, 1, seq_len)
+        x_reshaped = x.reshape(batch_size * num_channels, 1, seq_len)
+        
+        # Apply conv patching to all channels
+        patched = self.conv_patching(x_reshaped)  # (batch_size * num_channels, conv_embed_dim, num_patches)
+        
+        # Reshape back to separate batch and channel dimensions
+        _, conv_embed_dim, num_patches = patched.shape
+        patched = patched.reshape(batch_size, num_channels, conv_embed_dim, num_patches)
+        
+        # Transpose to get (batch_size, num_channels, num_patches, conv_embed_dim)
+        patched = patched.transpose(2, 3)
+        
+        return patched
+
+
+class DualRoPEAttention(nn.Module):
+    """Multi-head attention with separate RoPE for temporal and learnable RoPE for channels"""
+    def __init__(self,
+                 embed_dim: int = 256,
+                 num_heads: int = 8,
+                 dropout: float = 0.1,
+                 attention_type: str = "temporal",  # "temporal" or "channel"
+                 num_channels: int = 21,
+                 shared_channel_rope: Optional[nn.Module] = None):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.attention_type = attention_type
+        
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+        
+        # Linear projections
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        
+        # RoPE embeddings - different for temporal vs channel
+        if attention_type == "temporal":
+            # Standard RoPE for temporal attention
+            self.rotary_emb = RotaryEmbedding(
+                dim=self.head_dim,
+                theta=10000,
+                learned_freq=False
+            )
+        elif attention_type == "channel":
+            # Use shared learnable RoPE for channel attention if provided
+            if shared_channel_rope is not None:
+                self.rotary_emb = shared_channel_rope
+            else:
+                # Fallback to creating own RoPE
+                self.rotary_emb = RotaryEmbedding(
+                    dim=self.head_dim,
+                    theta=10000,
+                    learned_freq=True  # Learnable frequencies for channels
+                )
+        else:
+            raise ValueError(f"Unknown attention_type: {attention_type}")
+        
+        self.dropout = nn.Dropout(dropout)
+        self.scale = self.head_dim ** -0.5
+        
+    def forward(self, x, position_ids=None):
+        """
+        Args:
+            x: (batch_size, seq_len, embed_dim)
+            position_ids: (batch_size, seq_len) or (seq_len,) - custom position indices for RoPE
+        Returns:
+            (batch_size, seq_len, embed_dim)
+        """
+        batch_size, seq_len, embed_dim = x.shape
+        
+        # Linear projections
+        q = self.q_proj(x)
+        k = self.k_proj(x)
+        v = self.v_proj(x)
+        
+        # Reshape for multi-head attention
+        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        
+        # Apply RoPE
+        q = self.rotary_emb.rotate_queries_or_keys(q, position_ids=position_ids)
+        k = self.rotary_emb.rotate_queries_or_keys(k, position_ids=position_ids)
+        
+        # Scaled dot-product attention
+        attn_weights = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+        attn_weights = F.softmax(attn_weights, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+        
+        # Apply attention to values
+        attn_output = torch.matmul(attn_weights, v)
+        
+        # Reshape and project output
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)
+        output = self.out_proj(attn_output)
+        
+        return output
+
+
+class DualTransformerBlock(nn.Module):
+    """Biosignal transformer block with channel and temporal attention using dual RoPE"""
+    def __init__(self,
+                 embed_dim: int = 256,
+                 num_heads: int = 8,
+                 num_temporal_layers: int = 2,
+                 dropout: float = 0.1,
+                 mlp_ratio: float = 4.0,
+                 num_channels: int = 21,
+                 activation: str = "swiglu",
+                 norm_type: str = "rmsnorm",
+                 mlp_bias: bool = False,
+                 shared_channel_rope: Optional[nn.Module] = None):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_temporal_layers = num_temporal_layers
+        
+        # Helper function to create normalization layer
+        def create_norm(dim):
+            if norm_type == "rmsnorm":
+                return RMSNorm(dim)
+            elif norm_type == "layernorm":
+                return nn.LayerNorm(dim)
+            else:
+                raise ValueError(f"Unknown norm_type: {norm_type}")
+        
+        # Channel-wise attention with shared learnable RoPE
+        self.channel_attention = DualRoPEAttention(
+            embed_dim, num_heads, dropout, 
+            attention_type="channel", num_channels=num_channels,
+            shared_channel_rope=shared_channel_rope
+        )
+        self.channel_norm = create_norm(embed_dim)
+        
+        # Temporal attention layers with standard RoPE
+        self.temporal_attention_layers = nn.ModuleList([
+            DualRoPEAttention(embed_dim, num_heads, dropout, attention_type="temporal") 
+            for _ in range(num_temporal_layers)
+        ])
+        self.temporal_norms = nn.ModuleList([
+            create_norm(embed_dim) 
+            for _ in range(num_temporal_layers)
+        ])
+        
+        # MLP layers
+        mlp_hidden_dim = int(embed_dim * mlp_ratio)
+        self.channel_mlp = MLP(
+            dim=embed_dim,
+            hidden_dim=mlp_hidden_dim,
+            dropout=dropout,
+            activation=activation,
+            bias=mlp_bias
+        )
+        
+        self.temporal_mlps = nn.ModuleList([
+            MLP(
+                dim=embed_dim,
+                hidden_dim=mlp_hidden_dim,
+                dropout=dropout,
+                activation=activation,
+                bias=mlp_bias
+            ) for _ in range(num_temporal_layers)
+        ])
+        
+        self.channel_mlp_norm = create_norm(embed_dim)
+        self.temporal_mlp_norms = nn.ModuleList([
+            create_norm(embed_dim) 
+            for _ in range(num_temporal_layers)
+        ])
+        
+    def forward(self, x, temporal_position_ids=None):
+        """
+        Args:
+            x: (batch_size, num_channels, num_patches, embed_dim)
+            temporal_position_ids: (batch_size, num_patches) or (num_patches,) - position indices for temporal RoPE
+        Returns:
+            (batch_size, num_channels, num_patches, embed_dim)
+        """
+        batch_size, num_channels, num_patches, embed_dim = x.shape
+        
+        # 1. Channel-wise attention on each patch independently
+        x_for_channel_attn = x.permute(0, 2, 1, 3).contiguous().reshape(batch_size * num_patches, num_channels, embed_dim)
+        
+        # Apply channel attention with learnable RoPE
+        channel_attn_out = self.channel_attention(x_for_channel_attn)
+        
+        # Residual connection and layer norm
+        x_for_channel_attn = self.channel_norm(x_for_channel_attn + channel_attn_out)
+        
+        # MLP
+        channel_mlp_out = self.channel_mlp(x_for_channel_attn)
+        x_for_channel_attn = self.channel_mlp_norm(x_for_channel_attn + channel_mlp_out)
+        
+        # Reshape back
+        x = x_for_channel_attn.reshape(batch_size, num_patches, num_channels, embed_dim).permute(0, 2, 1, 3)
+        
+        # 2. Temporal attention on patches for each channel
+        x_for_temporal_attn = x.reshape(batch_size * num_channels, num_patches, embed_dim)
+        
+        # Prepare temporal position IDs
+        if temporal_position_ids is not None:
+            if temporal_position_ids.ndim == 2:
+                temporal_pos_ids_expanded = temporal_position_ids[0]
+            else:
+                temporal_pos_ids_expanded = temporal_position_ids
+        else:
+            temporal_pos_ids_expanded = None
+        
+        # Apply multiple temporal attention layers
+        for i in range(self.num_temporal_layers):
+            temporal_attn_out = self.temporal_attention_layers[i](x_for_temporal_attn, position_ids=temporal_pos_ids_expanded)
+            x_for_temporal_attn = self.temporal_norms[i](x_for_temporal_attn + temporal_attn_out)
+            
+            temporal_mlp_out = self.temporal_mlps[i](x_for_temporal_attn)
+            x_for_temporal_attn = self.temporal_mlp_norms[i](x_for_temporal_attn + temporal_mlp_out)
+        
+        # Reshape back
+        x = x_for_temporal_attn.reshape(batch_size, num_channels, num_patches, embed_dim)
+        
+        return x
+
+
+# ============================================================================
+# End of Pure Transformer Architecture Components
+# ============================================================================
+
+
+def _build_signal_tower(
+        embed_dim: int,
+        signal_cfg,
+        output_tokens: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    """Build a biosignals encoder tower
+    
+    Args:
+        embed_dim: Output embedding dimension
+        signal_cfg: BiosignalsCfg or dict with configuration
+        output_tokens: Whether to output tokens for multimodal decoder
+        cast_dtype: Optional dtype for casting
+    
+    Returns:
+        Biosignals encoder (either BiosignalsEncoder or PureTransformerBiosignalsEncoder)
+    """
+    if isinstance(signal_cfg, dict):
+        signal_cfg = BiosignalsCfg(**signal_cfg)
+    
+    import logging
+    architecture = getattr(signal_cfg, 'architecture', 'conv_transformer')
+    logging.info(f"Building biosignals encoder with architecture: {architecture}")
+    
+    if architecture == "pure_transformer":
+        signal_encoder = PureTransformerBiosignalsEncoder(
+            biosignals_cfg=signal_cfg,
+            embed_dim=embed_dim,
+            output_tokens=output_tokens,
+            cast_dtype=cast_dtype
+        )
+        logging.info(f"Pure Transformer architecture:")
+        logging.info(f"  Patch size: {signal_cfg.patch_size}")
+        logging.info(f"  Conv embed dim: {signal_cfg.conv_embed_dim}")
+        logging.info(f"  Transformer blocks: {signal_cfg.transformer_layers}")
+        logging.info(f"  Temporal layers per block: {signal_cfg.num_temporal_layers}")
+        logging.info(f"  Activation: {signal_cfg.activation}")
+        logging.info(f"  Norm type: {signal_cfg.norm_type}")
+        logging.info(f"  Share channel RoPE: {signal_cfg.share_channel_rope}")
+    elif architecture == "conv_transformer":
+        signal_encoder = BiosignalsEncoder(
+            biosignals_cfg=signal_cfg,
+            embed_dim=embed_dim,
+            output_tokens=output_tokens,
+            cast_dtype=cast_dtype
+        )
+        logging.info(f"Conv-Transformer architecture:")
+        logging.info(f"  Conv layers: {signal_cfg.conv_layers}")
+        logging.info(f"  Kernel sizes: {signal_cfg.kernel_sizes}")
+        logging.info(f"  Strides: {signal_cfg.strides}")
+        logging.info(f"  Transformer layers: {signal_cfg.transformer_layers}")
+    else:
+        raise ValueError(f"Unknown architecture: {architecture}. Must be 'conv_transformer' or 'pure_transformer'")
+    
+    return signal_encoder
+
+
+def _build_text_decoder_tower_v2(
+        embed_dim,
+        multimodal_cfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+        decoder_type: str = "cross_attention",
+        prefix_len: int = 0,
+):
+    """Build text decoder tower with support for different decoder types.
+    
+    Args:
+        embed_dim: Embedding dimension
+        multimodal_cfg: MultimodalCfg config
+        quick_gelu: Whether to use QuickGELU
+        cast_dtype: Optional dtype for casting
+        decoder_type: "cross_attention" or "concat"
+            - "cross_attention": Uses separate cross-attention layers (default CoCa)
+            - "concat": Concatenates image/biosignals and text tokens
+        prefix_len: Number of prefix tokens (condition embeddings) prepended to text
+            Used to pre-build prefix-causal attention mask
+    """
+    multimodal_cfg = MultimodalCfg(**multimodal_cfg) if isinstance(multimodal_cfg, dict) else multimodal_cfg
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+    norm_layer = (
+        LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+    )
+
+    if decoder_type == "cross_attention":
+        decoder = MultimodalTransformer(
+            context_length=multimodal_cfg.context_length,
+            width=multimodal_cfg.width,
+            heads=multimodal_cfg.heads,
+            layers=multimodal_cfg.layers,
+            mlp_ratio=multimodal_cfg.mlp_ratio,
+            ls_init_value=multimodal_cfg.ls_init_value,
+            output_dim=embed_dim,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            prefix_len=prefix_len,
+        )
+    elif decoder_type == "concat":
+        decoder = ConcatMultimodalTransformer(
+            context_length=multimodal_cfg.context_length,
+            width=multimodal_cfg.width,
+            heads=multimodal_cfg.heads,
+            layers=multimodal_cfg.layers,
+            mlp_ratio=multimodal_cfg.mlp_ratio,
+            ls_init_value=multimodal_cfg.ls_init_value,
+            output_dim=embed_dim,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            prefix_len=prefix_len,
+        )
+    else:
+        raise ValueError(f"Unknown decoder_type: {decoder_type}. Must be 'cross_attention' or 'concat'")
+
+    return decoder
+
+
+@dataclass
+class BiosignalsCfg:
+    """Configuration for biosignals encoder"""
+    input_channels: int = 12  # Number of input channels (e.g., 12-lead ECG)
+    signal_length: int = 1000  # Length of input time series
+    sampling_rate: int = 500  # Sampling rate in Hz
+    
+    # Architecture selection
+    architecture: str = "conv_transformer"  # "conv_transformer" or "pure_transformer"
+    
+    # Architecture parameters for conv_transformer
+    conv_layers: List[int] = None  # Conv layer dimensions
+    kernel_sizes: List[int] = None  # Kernel sizes for conv layers
+    strides: List[int] = None  # Strides for conv layers
+    
+    # Architecture parameters for pure_transformer
+    patch_size: int = 32  # Patch size for pure_transformer
+    conv_embed_dim: int = 256  # Conv embedding dimension for pure_transformer
+    num_temporal_layers: int = 2  # Number of temporal attention layers per block
+    activation: str = "swiglu"  # "swiglu", "gelu", "relu" (for pure_transformer)
+    norm_type: str = "rmsnorm"  # "rmsnorm", "layernorm" (for pure_transformer)
+    mlp_bias: bool = False  # Whether to use bias in MLP layers (for pure_transformer)
+    share_channel_rope: bool = True  # Share channel RoPE across blocks (for pure_transformer)
+    decoder_tokens: int = 32  # Number of decoder tokens for dual-axis transformer (pure_transformer)
+    
+    # Transformer parameters (shared)
+    transformer_layers: int = 6  # Number of transformer layers/blocks
+    transformer_width: int = 768  # Transformer width
+    transformer_heads: int = 12  # Number of attention heads
+    mlp_ratio: float = 4.0  # MLP expansion ratio
+    
+    # Pooling and output
+    pool_type: str = 'attn'  # 'avg', 'max', 'cls', 'attn'
+    dropout: float = 0.1
+    
+    def __post_init__(self):
+        if self.architecture == "conv_transformer":
+            if self.conv_layers is None:
+                # Default conv layers for processing time series
+                self.conv_layers = [64, 128, 256, 512]
+            if self.kernel_sizes is None:
+                # Default kernel sizes
+                self.kernel_sizes = [7, 5, 3, 3]
+            if self.strides is None:
+                # Default strides
+                self.strides = [2, 2, 2, 2]
+
+
+class BaseBiosignalsEncoder(nn.Module):
+    """
+    Base class for biosignals encoders that handles common pooling and projection logic.
+    Child classes should implement _encode() to return features before pooling.
+    """
+    
+    def __init__(
+        self,
+        biosignals_cfg: BiosignalsCfg,
+        embed_dim: int,
+        output_tokens: bool,
+        transformer_width: int,
+        cast_dtype: Optional[torch.dtype] = None
+    ):
+        super().__init__()
+        self.biosignals_cfg = biosignals_cfg
+        self.embed_dim = embed_dim
+        self.output_tokens = output_tokens
+        self.transformer_width = transformer_width
+        self.pool_type = biosignals_cfg.pool_type
+        
+        # Projection to output embedding dimension
+        self.proj_to_embed = nn.Linear(transformer_width, embed_dim)
+        
+        # Attention pooling if needed
+        if self.pool_type == 'attn':
+            self.attn_pool = nn.MultiheadAttention(
+                transformer_width,
+                biosignals_cfg.transformer_heads,
+                batch_first=True
+            )
+        
+    def _pool_features(self, x: torch.Tensor, has_cls_token: bool) -> torch.Tensor:
+        """
+        Pool features using the configured pooling method.
+        
+        Args:
+            x: Features of shape (batch_size, seq_len, width)
+            has_cls_token: Whether the sequence includes a CLS token at the last position
+            
+        Returns:
+            pooled: Pooled features of shape (batch_size, width)
+        """
+        if self.pool_type == 'cls':
+            # Use class token (last position)
+            pooled = x[:, -1]
+        elif self.pool_type == 'avg':
+            # Average pooling over sequence
+            if has_cls_token:
+                pooled = x[:, :-1].mean(dim=1)
+            else:
+                pooled = x.mean(dim=1)
+        elif self.pool_type == 'max':
+            # Max pooling over sequence
+            if has_cls_token:
+                pooled = x[:, :-1].max(dim=1)[0]
+            else:
+                pooled = x.max(dim=1)[0]
+        elif self.pool_type == 'attn':
+            # Attention pooling using cls token as query
+            query = x[:, -1:]  # CLS token as query
+            # CLS attends to content tokens
+            pooled, _ = self.attn_pool(query, x[:, :-1], x[:, :-1])
+            pooled = pooled.squeeze(1)
+        else:
+            raise ValueError(f"Unknown pool_type: {self.pool_type}")
+        
+        return pooled
+    
+    def _encode(self, biosignals: torch.Tensor) -> Tuple[torch.Tensor, bool]:
+        """
+        Encode biosignals to features. Must be implemented by child classes.
+        
+        Args:
+            biosignals: Input biosignals tensor
+            
+        Returns:
+            features: Encoded features of shape (batch_size, seq_len, transformer_width)
+            has_cls_token: Whether the sequence includes a CLS token at the last position
+        """
+        raise NotImplementedError("Child classes must implement _encode()")
+    
+    def forward(self, biosignals: torch.Tensor):
+        """
+        Forward pass with encoding, pooling, and projection.
+        
+        Args:
+            biosignals: Input biosignals tensor
+            
+        Returns:
+            embedding: Global embedding (batch_size, embed_dim)
+            tokens_for_decoder: Optional tokens for decoder (batch_size, seq_len, transformer_width)
+        """
+        # Encode to features
+        features, has_cls_token = self._encode(biosignals)
+        
+        # Pool features
+        pooled = self._pool_features(features, has_cls_token)
+        
+        # Project to final embedding dimension
+        embedding = self.proj_to_embed(pooled)
+        
+        if self.output_tokens:
+            # Return tokens for multimodal decoder
+            if has_cls_token:
+                # Exclude CLS token from tokens for decoder
+                tokens_for_decoder = features[:, :-1]
+            else:
+                tokens_for_decoder = features
+            return embedding, tokens_for_decoder
+        else:
+            return embedding
+    
+    def set_grad_checkpointing(self, enable=True):
+        # For compatibility with other models
+        pass
+
+
+class Conv1dBlock(nn.Module):
+    """1D Convolutional block with normalization and activation"""
+    
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, 
+                 norm_layer=nn.BatchNorm1d, act_layer=nn.ReLU):
+        super().__init__()
+        self.conv = nn.Conv1d(
+            in_channels, out_channels, kernel_size, 
+            stride=stride, padding=kernel_size//2
+        )
+        self.norm = norm_layer(out_channels)
+        self.act = act_layer()
+        self.dropout = nn.Dropout(0.1)
+        
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        x = self.act(x)
+        x = self.dropout(x)
+        return x
+
+
+class BiosignalsEncoder(BaseBiosignalsEncoder):
+    """
+    Biosignals encoder that converts time series data to embeddings.
+    Uses a combination of 1D convolutions and transformers.
+    """
+    
+    def __init__(
+        self, 
+        biosignals_cfg: BiosignalsCfg,
+        embed_dim: int = 512,
+        output_tokens: bool = False,
+        cast_dtype: Optional[torch.dtype] = None
+    ):
+        # Initialize base class with common pooling/projection logic
+        super().__init__(
+            biosignals_cfg=biosignals_cfg,
+            embed_dim=embed_dim,
+            output_tokens=output_tokens,
+            transformer_width=biosignals_cfg.transformer_width,
+            cast_dtype=cast_dtype
+        )
+        
+        # Convolutional feature extraction
+        conv_layers = []
+        in_channels = biosignals_cfg.input_channels
+        
+        for i, (out_channels, kernel_size, stride) in enumerate(
+            zip(biosignals_cfg.conv_layers, biosignals_cfg.kernel_sizes, biosignals_cfg.strides)
+        ):
+            conv_layers.append(
+                Conv1dBlock(in_channels, out_channels, kernel_size, stride)
+            )
+            in_channels = out_channels
+            
+        self.conv_layers = nn.Sequential(*conv_layers)
+        
+        # Calculate the length after convolutions with padding - we'll use a dummy forward pass
+        # to get the exact dimensions
+        with torch.no_grad():
+            dummy_input = torch.randn(1, biosignals_cfg.input_channels, biosignals_cfg.signal_length)
+            dummy_output = self.conv_layers(dummy_input)
+            conv_output_length = dummy_output.shape[2]
+        
+        self.conv_output_length = conv_output_length
+        self.conv_output_dim = biosignals_cfg.conv_layers[-1]
+        
+        # Projection to transformer dimension
+        self.proj_conv_to_transformer = nn.Linear(
+            self.conv_output_dim, biosignals_cfg.transformer_width
+        )
+        
+        # Positional embeddings for sequence positions (excluding CLS token)
+        # CLS token gets no positional embedding as it represents global context
+        self.pos_embed = nn.Parameter(
+            torch.randn(1, conv_output_length, biosignals_cfg.transformer_width)
+        )
+        
+        # Add a class token for global representation (only used for 'cls' and 'attn' pooling)
+        self.cls_token = nn.Parameter(
+            torch.randn(1, 1, biosignals_cfg.transformer_width)
+        )
+        
+        # Transformer layers
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+        act_layer = QuickGELU
+        
+        self.transformer_layers = nn.ModuleList([
+            TransformerBlock(
+                biosignals_cfg.transformer_width,
+                biosignals_cfg.transformer_heads,
+                biosignals_cfg.mlp_ratio,
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                dropout=biosignals_cfg.dropout
+            )
+            for _ in range(biosignals_cfg.transformer_layers)
+        ])
+        
+        # Final layer norm
+        self.ln_final = norm_layer(biosignals_cfg.transformer_width)
+            
+    def _encode(self, biosignals):
+        """
+        Encode biosignals to features before pooling.
+        
+        Args:
+            biosignals: Tensor of shape (batch_size, channels, signal_length)
+        Returns:
+            features: Encoded features of shape (batch_size, seq_len, transformer_width)
+            has_cls_token: Whether the sequence includes a CLS token at the last position
+        """
+        batch_size = biosignals.shape[0]
+        
+        # Apply convolutional layers
+        x = self.conv_layers(biosignals)  # (batch_size, conv_dim, conv_length)
+        
+        # Transpose to (batch_size, conv_length, conv_dim)
+        x = x.transpose(1, 2)
+        
+        # Project to transformer dimension
+        x = self.proj_conv_to_transformer(x)  # (batch_size, conv_length, transformer_width)
+        
+        # Add positional embeddings
+        x = x + self.pos_embed
+        
+        # Add class token only if needed for pooling
+        # For consistency with causal text encoder, append CLS token (not prepend)
+        if self.pool_type in ['cls', 'attn']:
+            cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+            x = torch.cat([x, cls_tokens], dim=1)  # (batch_size, conv_length + 1, transformer_width)
+            has_cls_token = True
+        else:
+            has_cls_token = False
+        
+        # Apply transformer layers
+        for layer in self.transformer_layers:
+            x = layer(x)
+            
+        # Apply final layer norm
+        x = self.ln_final(x)
+        
+        return x, has_cls_token
+
+
+class TransformerBlock(nn.Module):
+    """Transformer block with self-attention and MLP"""
+    
+    def __init__(
+        self, 
+        width: int, 
+        heads: int, 
+        mlp_ratio: float = 4.0,
+        act_layer=QuickGELU,
+        norm_layer=LayerNorm,
+        dropout: float = 0.1
+    ):
+        super().__init__()
+        self.attention = nn.MultiheadAttention(width, heads, dropout=dropout, batch_first=True)
+        self.ln_1 = norm_layer(width)
+        self.mlp = nn.Sequential(
+            nn.Linear(width, int(width * mlp_ratio)),
+            act_layer(),
+            nn.Dropout(dropout),
+            nn.Linear(int(width * mlp_ratio), width),
+            nn.Dropout(dropout)
+        )
+        self.ln_2 = norm_layer(width)
+        
+    def forward(self, x):
+        # Self-attention
+        attn_out, _ = self.attention(x, x, x)
+        x = x + attn_out
+        x = self.ln_1(x)
+        
+        # MLP
+        mlp_out = self.mlp(x)
+        x = x + mlp_out
+        x = self.ln_2(x)
+        
+        return x
+
+
+class AttnPooler(nn.Module):
+    """
+    CoCa-style attentional pooler.
+    A small multi-head attention layer with n_query learned queries (Q),
+    and the encoder sequence as both K and V. This lets us:
+      - n_query = 1  => global embedding for contrastive loss
+      - n_query = N  => compressed token set for decoder cross-attention
+    Ref: CoCa uses task-specific attentional pooling with nquery=1 for contrastive
+    and nquery=256 for generative objectives.  [oai_citation:2‡Medium](https://medium.com/%40arithmancylabs/coca-contrastive-captioners-are-image-textfoundation-models-324022377630?utm_source=chatgpt.com)
+    """
+    def __init__(self, dim: int, num_heads: int, n_query: int):
+        super().__init__()
+        self.n_query = n_query
+        self.query_tokens = nn.Parameter(torch.randn(1, n_query, dim) * 0.02)
+        self.attn = nn.MultiheadAttention(
+            embed_dim=dim,
+            num_heads=num_heads,
+            batch_first=True
+        )
+
+    def forward(self, x_seq: torch.Tensor) -> torch.Tensor:
+        """
+        x_seq: (B, L, D)
+        returns:
+            pooled: (B, n_query, D)
+        """
+        B = x_seq.size(0)
+        q = self.query_tokens.expand(B, -1, -1)  # (B, n_query, D)
+        pooled, _ = self.attn(q, x_seq, x_seq)   # pooled attends over all tokens
+        return pooled  # (B, n_query, D)
+
+
+class PureTransformerBiosignalsEncoder(BaseBiosignalsEncoder):
+    """
+    Pure Transformer encoder for biosignals with channel+temporal attention.
+
+    Updated to use CoCa-style task-specific attentional pooling:
+    - contrastive_pooler (n_query=1) → 1 global token for contrastive / CLS
+    - decoder_pooler (n_query=N_dec) → small set of summary tokens for text decoder
+
+    We still:
+      1. Patch each channel independently
+      2. Alternate channel-attn and temporal-attn in DualTransformerBlocks (factorized attention)
+      3. Keep (B, C, T, D) internally (cheap attention along channel or time separately)
+      4. Flatten to (B, C*T, D) only at the end
+      5. Run two poolers:
+          - 1-query pooler -> global token
+          - multi-query pooler -> decoder tokens
+      6. Append the 1-query pooled token to the end of x_seq so BaseBiosignalsEncoder
+         can keep using pool_type='cls' or 'attn' the same way.
+      7. Save the multi-query pooled tokens so, when output_tokens=True, we can hand
+         them to the text decoder instead of the full ~C*T sequence.
+
+    This mirrors CoCa's "task-specific attentional pooling," where the same encoder
+    supports both contrastive global alignment and caption-style generation with
+    minimal extra cost.  [oai_citation:3‡Medium](https://medium.com/%40arithmancylabs/coca-contrastive-captioners-are-image-textfoundation-models-324022377630?utm_source=chatgpt.com)
+    """
+
+    def __init__(
+        self,
+        biosignals_cfg: BiosignalsCfg,
+        embed_dim: int = 512,
+        output_tokens: bool = False,
+        cast_dtype: Optional[torch.dtype] = None
+    ):
+        super().__init__(
+            biosignals_cfg=biosignals_cfg,
+            embed_dim=embed_dim,
+            output_tokens=output_tokens,
+            transformer_width=biosignals_cfg.transformer_width,
+            cast_dtype=cast_dtype
+        )
+
+        # --- Sanity checks for RoPE dimensions ---
+        assert biosignals_cfg.transformer_width % biosignals_cfg.transformer_heads == 0, (
+            f"transformer_width ({biosignals_cfg.transformer_width}) must be divisible by "
+            f"transformer_heads ({biosignals_cfg.transformer_heads})"
+        )
+        head_dim = biosignals_cfg.transformer_width // biosignals_cfg.transformer_heads
+        assert head_dim % 2 == 0, (
+            f"head_dim ({head_dim}) must be even for RoPE. "
+            f"Got transformer_width={biosignals_cfg.transformer_width}, "
+            f"transformer_heads={biosignals_cfg.transformer_heads}"
+        )
+
+        # 1. Channel patching (Conv1d tokenizer per channel)
+        self.patching = ChannelPatching(
+            patch_size=biosignals_cfg.patch_size,
+            conv_embed_dim=biosignals_cfg.conv_embed_dim,
+            num_channels=biosignals_cfg.input_channels
+        )
+
+        # number of temporal patches per channel
+        self.num_patches = biosignals_cfg.signal_length // biosignals_cfg.patch_size
+
+        # 2. Project patch embeddings to transformer_width
+        self.embed_projection = nn.Linear(
+            biosignals_cfg.conv_embed_dim,
+            biosignals_cfg.transformer_width
+        )
+
+        # 2a. Channel ID embedding (categorical channel identity)
+        self.channel_id_embed = nn.Embedding(
+            num_embeddings=biosignals_cfg.input_channels,
+            embedding_dim=biosignals_cfg.transformer_width,
+        )
+
+        # 3. Shared learnable RoPE for channel attention (optional)
+        if biosignals_cfg.share_channel_rope:
+            shared_head_dim = biosignals_cfg.transformer_width // biosignals_cfg.transformer_heads
+            self.shared_channel_rope = RotaryEmbedding(
+                dim=shared_head_dim,
+                theta=10000,
+                learned_freq=True  # learnable for channel axis
+            )
+        else:
+            self.shared_channel_rope = None
+
+        # 4. Dual-axis Transformer blocks (channel attention + temporal attention)
+        self.transformer_blocks = nn.ModuleList([
+            DualTransformerBlock(
+                embed_dim=biosignals_cfg.transformer_width,
+                num_heads=biosignals_cfg.transformer_heads,
+                num_temporal_layers=biosignals_cfg.num_temporal_layers,
+                dropout=biosignals_cfg.dropout,
+                mlp_ratio=biosignals_cfg.mlp_ratio,
+                num_channels=biosignals_cfg.input_channels,
+                activation=biosignals_cfg.activation,
+                norm_type=biosignals_cfg.norm_type,
+                mlp_bias=biosignals_cfg.mlp_bias,
+                shared_channel_rope=self.shared_channel_rope if biosignals_cfg.share_channel_rope else None
+            ) for _ in range(biosignals_cfg.transformer_layers)
+        ])
+
+        # 5. Final norm
+        norm_layer = (
+            LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+        )
+        if biosignals_cfg.norm_type == "rmsnorm":
+            self.ln_final = RMSNorm(biosignals_cfg.transformer_width)
+        else:
+            self.ln_final = norm_layer(biosignals_cfg.transformer_width)
+
+        # 6. CoCa-style attentional poolers
+        #    - contrastive_pooler: n_query = 1 for global CLS token (contrastive head)
+        #    - decoder_pooler: n_query = decoder_tokens (e.g. 32) for compressed memory
+        #
+        # We'll add a new config field on BiosignalsCfg: decoder_tokens (int, default 32).
+        n_decoder_tokens = getattr(biosignals_cfg, "decoder_tokens", 32)
+
+        self.contrastive_pooler = AttnPooler(
+            dim=biosignals_cfg.transformer_width,
+            num_heads=biosignals_cfg.transformer_heads,
+            n_query=1
+        )
+
+        self.decoder_pooler = AttnPooler(
+            dim=biosignals_cfg.transformer_width,
+            num_heads=biosignals_cfg.transformer_heads,
+            n_query=n_decoder_tokens
+        )
+
+
+    def _encode(self, biosignals: torch.Tensor):
+        """
+        Returns:
+            features: (B, N_dec + 1, D)
+                first N_dec tokens  = pooled decoder tokens
+                last token          = global pooled token (contrastive CLS)
+            has_cls_token: True
+        """
+        B = biosignals.shape[0]
+        device = biosignals.device
+
+        # 1. Patch per channel -> (B, C, T, conv_dim)
+        x = self.patching(biosignals)
+
+        # 2. Project to model dim -> (B, C, T, D)
+        x = self.embed_projection(x)
+
+        # 2a. Add channel ID embedding
+        _, C, T, D = x.shape
+        channel_ids = torch.arange(C, device=device)              # (C,)
+        channel_bias = self.channel_id_embed(channel_ids)         # (C, D)
+        channel_bias = channel_bias.view(1, C, 1, D).expand(B, C, T, D)
+        x = x + channel_bias
+
+        # 3. Temporal RoPE positions
+        pos_ids = torch.arange(self.num_patches, device=device)   # (T,)
+
+        # 4. Dual-axis transformer blocks (channel-attn + temporal-attn)
+        for block in self.transformer_blocks:
+            x = block(x, temporal_position_ids=pos_ids)            # stays (B, C, T, D)
+
+        # 5. Final norm
+        x = self.ln_final(x)                                      # (B, C, T, D)
+
+        # 6. Flatten channels×time to a sequence for pooling (not for decoder!)
+        x_seq = x.reshape(B, C * T, D)                            # (B, L, D) with L = C*T
+
+        # 7. Task-specific attentional pooling (CoCa-style)
+        # contrastive_pooler: n_query=1  -> global_token (B,1,D)
+        # decoder_pooler:    n_query=Nd -> dec_tokens    (B,Nd,D)
+        global_token = self.contrastive_pooler(x_seq)             # (B, 1, D)
+        dec_tokens   = self.decoder_pooler(x_seq)                 # (B, N_dec, D)
+
+        # 8. Build final feature sequence:
+        #    [decoder tokens..., global token] so that:
+        #    - features[:, :-1] = dec_tokens (for decoder cross-attn)
+        #    - features[:, -1]  = global_token (for contrastive / CLS pooling)
+        features = torch.cat([dec_tokens, global_token], dim=1)   # (B, N_dec+1, D)
+
+        has_cls_token = True
+        return features, has_cls_token
+
+
+class SignalReconstructionDecoder(nn.Module):
+    """
+    Lightweight transformer decoder for signal reconstruction.
+    Uses 2-3 transformer encoder layers + final MLP to reconstruct biosignals.
+    Note: Uses TransformerEncoder (self-attention only) since we don't need cross-attention.
+    """
+    
+    def __init__(
+        self,
+        input_dim: int = 768,
+        num_layers: int = 2,
+        num_heads: int = 4,  # Reduced from 8 for efficiency
+        output_channels: int = 10,
+        output_length: int = 1920,
+    ):
+        super().__init__()
+        
+        # Transformer encoder layers (self-attention + FFN)
+        # Using 2x feedforward (instead of 4x) for lighter decoder
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=input_dim,
+            nhead=num_heads,
+            dim_feedforward=input_dim * 2,  # 1536 for input_dim=768
+            batch_first=True,
+            norm_first=True,
+        )
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
+        
+        # Final MLP to project to signal space
+        # Reduced intermediate dimension for efficiency
+        self.to_signal = nn.Sequential(
+            nn.Linear(input_dim, input_dim // 2),
+            nn.ReLU(),
+            nn.Linear(input_dim // 2, output_channels * output_length),
+        )
+        
+        self.output_channels = output_channels
+        self.output_length = output_length
+    
+    def forward(self, encoder_features):
+        """
+        Args:
+            encoder_features: (B, seq_len, input_dim) - unprojected encoder features
+        Returns:
+            reconstructed: (B, output_channels, output_length)
+        """
+        B = encoder_features.shape[0]
+        
+        # Self-attention on encoder features
+        decoded = self.transformer(encoder_features)  # (B, seq_len, dim)
+        
+        # Global average pooling
+        pooled = decoded.mean(dim=1)  # (B, dim)
+        
+        # Project to signal space
+        signal_flat = self.to_signal(pooled)  # (B, output_channels * output_length)
+        
+        # Reshape to signal format
+        signal = signal_flat.reshape(B, self.output_channels, self.output_length)
+        
+        return signal
+
+
+class BiosignalsCoCa(nn.Module):
+    """
+    CoCa model adapted for biosignals-text contrastive learning.
+    Replaces the vision tower with a biosignals encoder.
+    
+    Supports two decoder types:
+        - "cross_attention": Separate cross-attention between text and biosignals (default CoCa)
+        - "concat": Concatenate biosignals and text tokens with prefix-causal masking
+    """
+    
+    def __init__(
+            self,
+            embed_dim,
+            multimodal_cfg: MultimodalCfg,
+            text_cfg: CLIPTextCfg,
+            biosignals_cfg: BiosignalsCfg,
+            quick_gelu: bool = False,
+            init_logit_scale: float = np.log(1 / 0.07),
+            init_logit_bias: Optional[float] = None,
+            nonscalar_logit_scale: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            pad_id: int = 0,
+            decoder_type: str = "cross_attention",
+            num_caption_channels: int = 12,  # Number of channel/modality embeddings (22 for channels, 4 for modalities)
+            prefix_len: int = 0,
+            use_signal_decoder: bool = False,  # NEW: Enable signal reconstruction
+    ):
+        super().__init__()
+        multimodal_cfg = MultimodalCfg(**multimodal_cfg) if isinstance(multimodal_cfg, dict) else multimodal_cfg
+        text_cfg = CLIPTextCfg(**text_cfg) if isinstance(text_cfg, dict) else text_cfg
+        biosignals_cfg = BiosignalsCfg(**biosignals_cfg) if isinstance(biosignals_cfg, dict) else biosignals_cfg
+
+        self.decoder_type = decoder_type
+        self.num_channels = num_caption_channels
+        self.use_signal_decoder = use_signal_decoder
+        
+        # Debug logging for channel configuration
+        import logging
+        logging.info(f"BiosignalsCoCa initialized with num_caption_channels={num_caption_channels}, prefix_len={prefix_len}")
+        if use_signal_decoder:
+            logging.info(f"Signal reconstruction decoder enabled")
+
+        self.text = _build_text_tower(
+            embed_dim=embed_dim,
+            text_cfg=text_cfg,
+            quick_gelu=quick_gelu,
+            cast_dtype=cast_dtype,
+        )
+
+        vocab_size = (
+            self.text.vocab_size  # for hf models
+            if hasattr(text_cfg, "hf_model_name") and text_cfg.hf_model_name is not None
+            else text_cfg.vocab_size
+        )
+        
+        # Replace visual tower with biosignals tower
+        self.biosignals = _build_signal_tower(
+            embed_dim=embed_dim,
+            signal_cfg=biosignals_cfg,
+            output_tokens=True,  # Need tokens for multimodal decoder
+            cast_dtype=cast_dtype,
+        )
+
+        self.text_decoder = _build_text_decoder_tower_v2(
+            vocab_size,
+            multimodal_cfg=multimodal_cfg,
+            quick_gelu=quick_gelu,
+            cast_dtype=cast_dtype,
+            decoder_type=decoder_type,
+            prefix_len=prefix_len,
+        )
+
+        lshape = [1] if nonscalar_logit_scale else []
+        self.logit_scale = nn.Parameter(torch.ones(lshape) * init_logit_scale)
+        if init_logit_bias is not None:
+            self.logit_bias = nn.Parameter(torch.ones(lshape) * init_logit_bias)
+        else:
+            self.logit_bias = None
+        self.pad_id = pad_id
+
+        self.context_length = multimodal_cfg.context_length
+        
+        # Learnable channel/modality embeddings
+        # num_caption_channels will be 23 for individual channel mode or 5 for modality mode
+        # Dimension should match the decoder width (multimodal_cfg.width for text decoder input)
+        self.channel_embeddings = nn.Parameter(
+            torch.randn(num_caption_channels, multimodal_cfg.width) * 0.02
+        )
+        
+        # Learnable padding embedding for -1 positions
+        # This learns to be "neutral" or ignored during training (similar to [PAD] tokens)
+        self.padding_embedding = nn.Parameter(
+            torch.randn(multimodal_cfg.width) * 0.02
+        )
+        
+        self.decoder_width = multimodal_cfg.width
+        
+        # Optional signal reconstruction decoder
+        if use_signal_decoder:
+            self.signal_decoder = SignalReconstructionDecoder(
+                input_dim=biosignals_cfg.transformer_width,
+                num_layers=2,  # Lightweight: 2 transformer layers
+                num_heads=biosignals_cfg.transformer_heads,
+                output_channels=biosignals_cfg.input_channels,
+                output_length=biosignals_cfg.signal_length,
+            )
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable: bool = True):
+        self.biosignals.set_grad_checkpointing(enable)
+        self.text.set_grad_checkpointing(enable)
+        self.text_decoder.set_grad_checkpointing(enable)
+
+    def lock_text_tower(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True):
+        """Lock the text encoder, optionally leaving the last N layers unlocked.
+        
+        Args:
+            unlocked_layers: Number of layers to leave unlocked (from the end)
+            freeze_layer_norm: Whether to freeze LayerNorm parameters in locked layers
+        """
+        if hasattr(self.text, 'lock'):
+            # For HFTextEncoder (Pythia, etc.)
+            self.text.lock(unlocked_layers, freeze_layer_norm)
+            
+            # IMPORTANT: Unfreeze newly added token embeddings (e.g., <pad>, <coca_cls>)
+            # These were randomly initialized and need to be trained
+            if hasattr(self.text, 'original_vocab_size'):
+                import logging
+                embedding_module = self.text.transformer.get_input_embeddings()
+                original_size = self.text.original_vocab_size
+                current_size = embedding_module.weight.shape[0]
+                
+                if current_size > original_size:
+                    # Enable gradients for the embedding layer
+                    embedding_module.weight.requires_grad = True
+                    
+                    # Store metadata for optimizer configuration (zero weight decay)
+                    self.text._new_token_start_idx = original_size
+                    
+                    # Get actual embedding size (may be padded for Tensor Cores)
+                    actual_embedding_size = embedding_module.weight.shape[0]
+                    new_vocab_size = self.text.vocab_size  # Actual number of tokens (not padded)
+                    
+                    # Register parameter-level hook to mask frozen token gradients
+                    # IMPORTANT: This is registered BEFORE DDP wrapping to ensure it persists
+                    def _zero_grad_frozen_tokens(grad):
+                        """Zero out gradients for old (frozen) tokens and padding, keep only new tokens."""
+                        if grad is not None:
+                            # Zero out pretrained tokens [0:original_size]
+                            grad[:original_size] = 0
+                            # Zero out padding tokens [new_vocab_size:actual_embedding_size]
+                            if actual_embedding_size > new_vocab_size:
+                                grad[new_vocab_size:] = 0
+                        return grad
+                    
+                    embedding_module.weight.register_hook(_zero_grad_frozen_tokens)
+                    
+                    num_new_tokens = new_vocab_size - original_size
+                    num_padding_tokens = actual_embedding_size - new_vocab_size
+                    logging.info(f"Embedding layer configuration:")
+                    logging.info(f"  Trainable new tokens: {num_new_tokens} (indices {original_size}:{new_vocab_size})")
+                    logging.info(f"  Frozen pretrained tokens: {original_size} (indices 0:{original_size})")
+                    if num_padding_tokens > 0:
+                        logging.info(f"  Frozen padding tokens: {num_padding_tokens} (indices {new_vocab_size}:{actual_embedding_size})")
+                    logging.info(f"  Total embedding size: {actual_embedding_size}")
+                    logging.info(f"Registered gradient masking hook before DDP wrapping")
+                    logging.info(f"NOTE: Optimizer uses weight_decay=0 for embedding layer")
+        else:
+            # For standard TextTransformer
+            assert False, "BiosignalsCoCa does not support locking standard TextTransformer"
+            from .transformer import lock_text_tower
+            lock_text_tower(self, unlocked_layers)
+
+    def _encode_biosignals(self, biosignals, normalize: bool = True):
+        biosignals_latent, tokens_embs = self.biosignals(biosignals)
+        biosignals_latent = F.normalize(biosignals_latent, dim=-1) if normalize else biosignals_latent
+        return biosignals_latent, tokens_embs
+
+    def _encode_text(self, text, normalize: bool = True):
+        text_latent, token_emb = self.text(text)
+        text_latent = F.normalize(text_latent, dim=-1) if normalize else text_latent
+        return text_latent, token_emb
+
+    def encode_image(self, biosignals, normalize: bool = True):
+        biosignals_latent, _ = self._encode_biosignals(biosignals, normalize=normalize)
+        return biosignals_latent
+
+    def encode_text(self, text, normalize: bool = True):
+        text_latent, _ = self._encode_text(text, normalize=normalize)
+        return text_latent
+
+    def _get_channel_condition_embs(self, channel_indices: torch.Tensor) -> torch.Tensor:
+        """Convert channel/modality indices to embeddings with learnable padding.
+        
+        Args:
+            channel_indices: (batch_size, prefix_len) tensor of indices
+                - Individual mode: indices into 23 channel embeddings (22 channels + 1 stage_event)
+                - Modality mode: indices into 5 modality embeddings (4 modalities + 1 stage_event)
+                - Padded with -1 for variable length (uses learnable padding_embedding for -1)
+            
+        Returns:
+            condition_embs: (batch_size, prefix_len, decoder_width)
+                Embeddings for all positions. -1 positions use learnable padding_embedding
+                that learns to be neutral/ignored during training.
+        """
+        batch_size, prefix_len = channel_indices.shape
+        
+        # Create output tensor
+        condition_embs = torch.zeros(batch_size, prefix_len, self.decoder_width, 
+                                     dtype=self.channel_embeddings.dtype, 
+                                     device=self.channel_embeddings.device)
+        
+        # Create mask for valid (non-padding) indices
+        valid_mask = channel_indices >= 0  # (batch_size, prefix_len)
+        padding_mask = channel_indices == -1  # (batch_size, prefix_len)
+        
+        # Gather channel embeddings for valid indices
+        # Clamp to 0 for safe indexing (will be overwritten by padding where needed)
+        indices_safe = channel_indices.clamp(min=0)
+        
+        # Expand embeddings for batching
+        expanded_embeddings = self.channel_embeddings.unsqueeze(0).expand(batch_size, -1, -1)
+        
+        # Gather embeddings
+        indices_expanded = indices_safe.unsqueeze(-1).expand(-1, -1, self.decoder_width)
+        gathered_embs = torch.gather(expanded_embeddings, 1, indices_expanded)
+        
+        # Fill in valid positions with gathered embeddings
+        condition_embs[valid_mask] = gathered_embs[valid_mask]
+        
+        # Fill in padding positions with learnable padding embedding
+        if padding_mask.any():
+            # Broadcast padding_embedding to all padding positions
+            condition_embs[padding_mask] = self.padding_embedding
+        
+        return condition_embs
+    
+    def forward(
+            self,
+            biosignals,
+            text: Optional[torch.Tensor] = None,
+            biosignals_latent: Optional[torch.Tensor] = None,
+            biosignals_embs: Optional[torch.Tensor] = None,
+
+            channel_indices: Optional[torch.Tensor] = None,
+            output_labels: bool = True,
+    ):
+        """Forward pass for BiosignalsCoCa model.
+        
+        Args:
+            biosignals: Input biosignals tensor
+            text: Optional text token ids
+            biosignals_latent: Optional pre-computed biosignals latent features
+            biosignals_embs: Optional pre-computed biosignals token embeddings
+
+            channel_indices: Optional (batch_size, num_selected_channels) tensor of channel indices
+                Used to select channel-specific condition embeddings. If provided, overrides condition_embs.
+            output_labels: Whether to output labels for loss computation
+        """
+        if biosignals_latent is None or biosignals_embs is None:
+            biosignals_latent, biosignals_embs = self._encode_biosignals(biosignals)
+
+        if text is None:
+            return {"image_features": biosignals_latent, "image_embs": biosignals_embs}
+
+        text_latent, token_embs = self._encode_text(text)
+
+        # FIXME this isn't an ideal solution, would like to improve -RW
+        labels: Optional[torch.Tensor] = text[:, 1:] if output_labels else None
+        if output_labels:
+            # align text_embs and thus logits with labels for teacher-forcing caption loss
+            token_embs = token_embs[:, :-1]
+        
+        # Convert channel indices to condition embeddings if provided
+        if channel_indices is not None:
+            condition_embs = self._get_channel_condition_embs(channel_indices)
+        else:
+            condition_embs = None
+
+        logits = self.text_decoder(biosignals_embs, token_embs, condition_embs=condition_embs)
+        out_dict = {
+            "image_features": biosignals_latent,
+            "text_features": text_latent,
+            "logits": logits,
+            "logit_scale": self.logit_scale.exp()
+        }
+        if labels is not None:
+            out_dict["labels"] = labels
+        if self.logit_bias is not None:
+            out_dict["logit_bias"] = self.logit_bias
+        
+        # Optional signal reconstruction
+        if self.use_signal_decoder:
+            reconstructed_signal = self.signal_decoder(biosignals_embs)
+            out_dict["reconstructed_signal"] = reconstructed_signal
+            out_dict["original_signal"] = biosignals
+        
+        return out_dict
+
+    def generate(
+        self,
+        biosignals,
+        text=None,
+        seq_len=30,
+        max_seq_len=256,
+        temperature=1.,
+        generation_type="beam_search",
+        top_p=0.1,
+        top_k=1,
+        pad_token_id=None,
+        eos_token_id=None,
+        sot_token_id=None,
+        num_beams=6,
+        num_beam_groups=3,
+        min_seq_len=5,
+        stopping_criteria=None,
+        repetition_penalty=1.0,
+        fixed_output_length=False,
+        condition_embs=None,
+        channel_indices=None,
+    ):
+# taking many ideas and components from HuggingFace GenerationMixin
+        # https://huggingface.co/docs/transformers/main/en/main_classes/text_generation
+        assert _has_transformers, "Please install transformers for generate functionality. `pip install transformers`."
+        assert seq_len > min_seq_len, "seq_len must be larger than min_seq_len"
+        device = biosignals.device
+        
+        # Note: condition_embs parameter is for backward compatibility
+        # We pass channel_indices directly to forward(), which handles the conversion internally
+
+        with torch.no_grad():
+            sot_token_id = _token_to_tensor(sot_token_id, device=device)
+            eos_token_id = _token_to_tensor(eos_token_id, device=device)
+            pad_token_id = pad_token_id
+            logit_processor = LogitsProcessorList(
+                [
+                    MinLengthLogitsProcessor(min_seq_len, eos_token_id),
+                    RepetitionPenaltyLogitsProcessor(repetition_penalty),
+                ]
+            )
+
+            if stopping_criteria is None:
+                stopping_criteria = [MaxLengthCriteria(max_length=seq_len)]
+            stopping_criteria = StoppingCriteriaList(stopping_criteria)
+
+            if generation_type == "beam_search":
+                output = self._generate_beamsearch(
+                    biosignals_inputs=biosignals,
+                    pad_token_id=pad_token_id,
+                    eos_token_id=eos_token_id,
+                    sot_token_id=sot_token_id,
+                    num_beams=num_beams,
+                    num_beam_groups=num_beam_groups,
+                    min_seq_len=min_seq_len,
+                    stopping_criteria=stopping_criteria,
+                    logit_processor=logit_processor,
+                    channel_indices=channel_indices,
+                )
+                if fixed_output_length and output.shape[1] < seq_len:
+                    pad_len = seq_len - output.shape[1]
+                    return torch.cat((
+                            output,
+                            torch.ones(output.shape[0], pad_len, device=device, dtype=output.dtype) * pad_token_id
+                        ),
+                        dim=1
+                    )
+                return output
+
+            elif generation_type == "top_p":
+                logit_warper = GENERATION_TYPES[generation_type](top_p)
+            elif generation_type == "top_k":
+                logit_warper = GENERATION_TYPES[generation_type](top_k)
+            else:
+                raise ValueError(
+                    f"generation_type has to be one of "
+                    f"{'| ' + ' | '.join(list(GENERATION_TYPES.keys())) + ' |'}."
+                )
+
+            biosignals_latent, biosignals_embs = self._encode_biosignals(biosignals)
+
+            if text is None:
+                text = torch.ones((biosignals.shape[0], 1), device=device, dtype=torch.long) * sot_token_id
+
+            was_training = self.training
+            num_dims = len(text.shape)
+
+            if num_dims == 1:
+                text = text[None, :]
+
+            self.eval()
+            out = text
+
+            while True:
+                x = out[:, -max_seq_len:]
+                cur_len = x.shape[1]
+                logits = self(
+                    biosignals,
+                    x,
+                    biosignals_latent=biosignals_latent,
+                    biosignals_embs=biosignals_embs,
+                    channel_indices=channel_indices,
+                    output_labels=False,
+                )["logits"][:, -1]
+                mask = (out[:, -1] == eos_token_id) | (out[:, -1] == pad_token_id)
+                sample = torch.ones((out.shape[0], 1), device=device, dtype=torch.long) * pad_token_id
+
+                if mask.all():
+                    if not fixed_output_length:
+                        break
+                else:
+                    logits = logits[~mask, :]
+                    filtered_logits = logit_processor(x[~mask, :], logits)
+                    filtered_logits = logit_warper(x[~mask, :], filtered_logits)
+                    probs = F.softmax(filtered_logits / temperature, dim=-1)
+
+                    if (cur_len + 1 == seq_len):
+                        sample[~mask, :] = torch.ones((sum(~mask), 1), device=device, dtype=torch.long) * eos_token_id
+                    else:
+                        sample[~mask, :] = torch.multinomial(probs, 1)
+
+                out = torch.cat((out, sample), dim=-1)
+
+                cur_len += 1
+
+                if all(stopping_criteria(out, None)):
+                    break
+
+            if num_dims == 1:
+                out = out.squeeze(0)
+
+            self.train(was_training)
+            return out
+
+    def _generate_beamsearch(
+            self,
+            biosignals_inputs,
+            pad_token_id=None,
+            eos_token_id=None,
+            sot_token_id=None,
+            num_beams=6,
+            num_beam_groups=3,
+            min_seq_len=5,
+            stopping_criteria=None,
+            logit_processor=None,
+            logit_warper=None,
+            channel_indices=None,
+    ):
+        device = biosignals_inputs.device
+        batch_size = biosignals_inputs.shape[0]
+        biosignals_inputs = torch.repeat_interleave(biosignals_inputs, num_beams, dim=0)
+        biosignals_latent, biosignals_embs = self._encode_biosignals(biosignals_inputs)
+        
+        # Repeat channel indices for beam search if provided
+        # forward() will convert them to condition embeddings internally
+        if channel_indices is not None:
+            channel_indices = torch.repeat_interleave(channel_indices, num_beams, dim=0)
+
+        input_ids = torch.ones((batch_size * num_beams, 1), device=device, dtype=torch.long)
+        input_ids = input_ids * sot_token_id
+        beam_scorer = BeamSearchScorer(
+            batch_size=batch_size,
+            num_beams=num_beams,
+            device=device,
+            num_beam_groups=num_beam_groups,
+        )
+        # instantiate logits processors
+        logits_processor = (
+            LogitsProcessorList([MinLengthLogitsProcessor(min_seq_len, eos_token_id=eos_token_id)])
+            if logit_processor is None
+            else logit_processor
+        )
+
+        num_beams = beam_scorer.num_beams
+        num_beam_groups = beam_scorer.num_beam_groups
+        num_sub_beams = num_beams // num_beam_groups
+        batch_size = len(beam_scorer._beam_hyps) // num_beam_groups
+        batch_beam_size, cur_len = input_ids.shape
+        beam_indices = None
+
+        if num_beams * batch_size != batch_beam_size:
+            raise ValueError(
+                f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}."
+            )
+
+        beam_scores = torch.full((batch_size, num_beams), -1e9, dtype=torch.float, device=device)
+        # initialise score of first beam of each group with 0 and the rest with 1e-9. This ensures that the beams in
+        # the same group don't produce same tokens everytime.
+        beam_scores[:, ::num_sub_beams] = 0
+        beam_scores = beam_scores.view((batch_size * num_beams,))
+
+        while True:
+
+            # predicted tokens in cur_len step
+            current_tokens = torch.zeros(batch_size * num_beams, dtype=input_ids.dtype, device=device)
+
+            # indices which will form the beams in the next time step
+            reordering_indices = torch.zeros(batch_size * num_beams, dtype=torch.long, device=device)
+
+            # do one decoder step on all beams of all sentences in batch
+            model_inputs = prepare_inputs_for_generation(input_ids=input_ids, biosignals_inputs=biosignals_inputs)
+            outputs = self(
+                model_inputs['biosignals'],
+                model_inputs['text'],
+                biosignals_latent=biosignals_latent,
+                biosignals_embs=biosignals_embs,
+                channel_indices=channel_indices,
+                output_labels=False,
+            )
+
+            for beam_group_idx in range(num_beam_groups):
+                group_start_idx = beam_group_idx * num_sub_beams
+                group_end_idx = min(group_start_idx + num_sub_beams, num_beams)
+                group_size = group_end_idx - group_start_idx
+
+                # indices of beams of current group among all sentences in batch
+                batch_group_indices = []
+
+                for batch_idx in range(batch_size):
+                    batch_group_indices.extend(
+                        [batch_idx * num_beams + idx for idx in range(group_start_idx, group_end_idx)]
+                    )
+                group_input_ids = input_ids[batch_group_indices]
+
+                # select outputs of beams of currentg group only
+                next_token_logits = outputs['logits'][batch_group_indices, -1, :]
+                vocab_size = next_token_logits.shape[-1]
+
+                next_token_scores_processed = logits_processor(
+                    group_input_ids, next_token_logits, current_tokens=current_tokens, beam_group_idx=beam_group_idx
+                )
+                next_token_scores = next_token_scores_processed + beam_scores[batch_group_indices].unsqueeze(-1)
+                next_token_scores = next_token_scores.expand_as(next_token_scores_processed)
+
+                # reshape for beam search
+                next_token_scores = next_token_scores.view(batch_size, group_size * vocab_size)
+
+                next_token_scores, next_tokens = torch.topk(
+                    next_token_scores, 2 * group_size, dim=1, largest=True, sorted=True
+                )
+
+                next_indices = torch.div(next_tokens, vocab_size, rounding_mode="floor")
+                next_tokens = next_tokens % vocab_size
+
+                # stateless
+                process_beam_indices = sum(beam_indices, ()) if beam_indices is not None else None
+                beam_outputs = beam_scorer.process(
+                    group_input_ids,
+                    next_token_scores,
+                    next_tokens,
+                    next_indices,
+                    pad_token_id=pad_token_id,
+                    eos_token_id=eos_token_id,
+                    beam_indices=process_beam_indices,
+                    group_index=beam_group_idx,
+                )
+                beam_scores[batch_group_indices] = beam_outputs["next_beam_scores"]
+                beam_next_tokens = beam_outputs["next_beam_tokens"]
+                beam_idx = beam_outputs["next_beam_indices"]
+
+                input_ids[batch_group_indices] = group_input_ids[beam_idx]
+                group_input_ids = torch.cat([group_input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)
+                current_tokens[batch_group_indices] = group_input_ids[:, -1]
+
+                # (beam_idx // group_size) -> batch_idx
+                # (beam_idx % group_size) -> offset of idx inside the group
+                reordering_indices[batch_group_indices] = (
+                    num_beams * torch.div(beam_idx, group_size, rounding_mode="floor") + group_start_idx + (beam_idx % group_size)
+                )
+
+            input_ids = torch.cat([input_ids, current_tokens.unsqueeze(-1)], dim=-1)
+
+            # increase cur_len
+            cur_len = cur_len + 1
+            if beam_scorer.is_done or all(stopping_criteria(input_ids, None)):
+                break
+
+        final_beam_indices = sum(beam_indices, ()) if beam_indices is not None else None
+        sequence_outputs = beam_scorer.finalize(
+            input_ids,
+            beam_scores,
+            next_tokens,
+            next_indices,
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            max_length=stopping_criteria.max_length,
+            beam_indices=final_beam_indices,
+        )
+        return sequence_outputs['sequences']
+
+
+def prepare_inputs_for_generation(input_ids, biosignals_inputs, past=None, **kwargs):
+    if past:
+        input_ids = input_ids[:, -1].unsqueeze(-1)
+
+    attention_mask = kwargs.get("attention_mask", None)
+    position_ids = kwargs.get("position_ids", None)
+
+    if attention_mask is not None and position_ids is None:
+        # create position_ids on the fly for batch generation
+        position_ids = attention_mask.long().cumsum(-1) - 1
+        position_ids.masked_fill_(attention_mask == 0, 1)
+    else:
+        position_ids = None
+    return {
+        "text": input_ids,
+        "biosignals": biosignals_inputs,
+        "past_key_values": past,
+        "position_ids": position_ids,
+        "attention_mask": attention_mask,
+    }

src/open_clip/bpe_simple_vocab_16e6.txt.gz

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

src/open_clip/coca_model.py

@@ -0,0 +1,586 @@
+from typing import Dict, List, Optional, Union
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+import numpy as np
+from dataclasses import dataclass
+
+from .transformer import (
+    LayerNormFp32,
+    LayerNorm,
+    QuickGELU,
+    MultimodalTransformer,
+)
+from .model import CLIPTextCfg, _build_text_tower
+
+try:
+    from transformers import (
+        BeamSearchScorer,
+        LogitsProcessorList,
+        TopPLogitsWarper,
+        TopKLogitsWarper,
+        RepetitionPenaltyLogitsProcessor,
+        MinLengthLogitsProcessor,
+        MaxLengthCriteria,
+        StopStringCriteria,
+        EosTokenCriteria,
+        StoppingCriteriaList
+    )
+
+    GENERATION_TYPES = {
+        "top_k": TopKLogitsWarper,
+        "top_p": TopPLogitsWarper,
+        "beam_search": "beam_search"
+    }
+    _has_transformers = True
+except ImportError as e:
+    GENERATION_TYPES = {
+        "top_k": None,
+        "top_p": None,
+        "beam_search": "beam_search"
+    }
+    _has_transformers = False
+
+
+@dataclass
+class MultimodalCfg(CLIPTextCfg):
+    mlp_ratio: int = 4
+    dim_head: int = 64
+    heads: int = 8
+    n_queries: int = 256
+    attn_pooler_heads: int = 8
+
+
+def _build_text_decoder_tower(
+        embed_dim,
+        multimodal_cfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    multimodal_cfg = MultimodalCfg(**multimodal_cfg) if isinstance(multimodal_cfg, dict) else multimodal_cfg
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+    norm_layer = (
+        LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+    )
+
+    decoder = MultimodalTransformer(
+        context_length=multimodal_cfg.context_length,
+        width=multimodal_cfg.width,
+        heads=multimodal_cfg.heads,
+        layers=multimodal_cfg.layers,
+        ls_init_value=multimodal_cfg.ls_init_value,
+        output_dim=embed_dim,
+        act_layer=act_layer,
+        norm_layer=norm_layer,
+    )
+
+    return decoder
+
+
+def _token_to_tensor(token_id, device: str = "cpu") -> torch.Tensor:
+    if not isinstance(token_id, torch.Tensor):
+        if isinstance(token_id, int):
+            token_id = [token_id]
+        token_id = torch.tensor(token_id, device=device)
+    return token_id
+
+
+class CoCa(nn.Module):
+    def __init__(
+            self,
+            embed_dim,
+            multimodal_cfg: MultimodalCfg,
+            text_cfg: CLIPTextCfg,
+            vision_cfg=None,
+            quick_gelu: bool = False,
+            init_logit_scale: float = np.log(1 / 0.07),
+            init_logit_bias: Optional[float] = None,
+            nonscalar_logit_scale: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            pad_id: int = 0,
+    ):
+        super().__init__()
+        multimodal_cfg = MultimodalCfg(**multimodal_cfg) if isinstance(multimodal_cfg, dict) else multimodal_cfg
+        text_cfg = CLIPTextCfg(**text_cfg) if isinstance(text_cfg, dict) else text_cfg
+
+        self.text = _build_text_tower(
+            embed_dim=embed_dim,
+            text_cfg=text_cfg,
+            quick_gelu=quick_gelu,
+            cast_dtype=cast_dtype,
+        )
+
+        vocab_size = (
+            self.text.vocab_size
+            if hasattr(text_cfg, "hf_model_name") and text_cfg.hf_model_name is not None
+            else text_cfg.vocab_size
+        )
+
+        if vision_cfg is not None:
+            from .model import CLIPVisionCfg, _build_vision_tower
+            vision_cfg = CLIPVisionCfg(**vision_cfg) if isinstance(vision_cfg, dict) else vision_cfg
+            self.visual = _build_vision_tower(
+                embed_dim=embed_dim,
+                vision_cfg=vision_cfg,
+                quick_gelu=quick_gelu,
+                cast_dtype=cast_dtype,
+            )
+        else:
+            self.visual = None
+
+        self.text_decoder = _build_text_decoder_tower(
+            vocab_size,
+            multimodal_cfg=multimodal_cfg,
+            quick_gelu=quick_gelu,
+            cast_dtype=cast_dtype,
+        )
+
+        lshape = [1] if nonscalar_logit_scale else []
+        self.logit_scale = nn.Parameter(torch.ones(lshape) * init_logit_scale)
+        if init_logit_bias is not None:
+            self.logit_bias = nn.Parameter(torch.ones(lshape) * init_logit_bias)
+        else:
+            self.logit_bias = None
+        self.pad_id = pad_id
+
+        self.context_length = multimodal_cfg.context_length
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable: bool = True):
+        self.visual.set_grad_checkpointing(enable)
+        self.text.set_grad_checkpointing(enable)
+        self.text_decoder.set_grad_checkpointing(enable)
+
+    def _encode_image(self, images, normalize: bool = True):
+        image_latent, tokens_embs = self.visual(images)
+        image_latent = F.normalize(image_latent, dim=-1) if normalize else image_latent
+        return image_latent, tokens_embs
+
+    def _encode_text(self, text, normalize: bool = True):
+        text_latent, token_emb = self.text(text)
+        text_latent = F.normalize(text_latent, dim=-1) if normalize else text_latent
+        return text_latent, token_emb
+
+    def encode_image(self, images, normalize: bool = True):
+        image_latent, _ = self._encode_image(images, normalize=normalize)
+        return image_latent
+
+    def encode_text(self, text, normalize: bool = True):
+        text_latent, _ = self._encode_text(text, normalize=normalize)
+        return text_latent
+
+    def forward_intermediates(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+            image_indices: Optional[Union[int, List[int]]] = None,
+            text_indices: Optional[Union[int, List[int]]] = None,
+            stop_early: bool = False,
+            normalize: bool = True,
+            normalize_intermediates: bool = False,
+            intermediates_only: bool = False,
+            image_output_fmt: str = 'NCHW',
+            image_output_extra_tokens: bool = False,
+            text_output_fmt: str = 'NLC',
+            text_output_extra_tokens: bool = False,
+            output_logits: bool = False,
+            output_logit_scale_bias: bool = False,
+    ) -> Dict[str, Union[torch.Tensor, List[torch.Tensor]]]:
+        """ Forward features that returns intermediates.
+
+        Args:
+            image: Input image tensor
+            text: Input text tensor
+            image_indices: For image tower, Take last n blocks if int, all if None, select matching indices if sequence
+            text_indices: Take last n blocks if int, all if None, select matching indices if sequence
+            stop_early: Stop iterating over blocks when last desired intermediate hit
+            normalize: L2 Normalize final image and text features (if present)
+            normalize_intermediates: Apply final encoder norm layer to all intermediates (if possible)
+            intermediates_only: Only return intermediate features, do not return final features
+            image_output_fmt: Shape of intermediate image feature outputs
+            image_output_extra_tokens: Return both prefix and spatial intermediate tokens
+            text_output_fmt: Shape of intermediate text feature outputs
+            text_output_extra_tokens: Return both prefix and spatial intermediate tokens
+            output_logits: Include logits in output
+            output_logit_scale_bias: Include the logit scale bias in the output
+        Returns:
+
+        """
+        output = {}
+        if intermediates_only:
+            # intermediates only disables final feature normalization, and include logits
+            normalize = False
+            output_logits = False
+        if output_logits:
+            assert False, 'FIXME, needs implementing'
+
+        if image is not None:
+            image_output = self.visual.forward_intermediates(
+                image,
+                indices=image_indices,
+                stop_early=stop_early,
+                normalize_intermediates=normalize_intermediates,
+                intermediates_only=intermediates_only,
+                output_fmt=image_output_fmt,
+                output_extra_tokens=image_output_extra_tokens,
+            )
+            if normalize and "image_features" in image_output:
+                image_output["image_features"] = F.normalize(image_output["image_features"], dim=-1)
+            output.update(image_output)
+
+        if text is not None:
+            text_output = self.text.forward_intermediates(
+                text,
+                indices=text_indices,
+                stop_early=stop_early,
+                normalize_intermediates=normalize_intermediates,
+                intermediates_only=intermediates_only,
+                output_fmt=text_output_fmt,
+                output_extra_tokens=text_output_extra_tokens,
+            )
+            if normalize and "text_features" in text_output:
+                text_output["text_features"] = F.normalize(text_output["text_features"], dim=-1)
+            output.update(text_output)
+
+        # FIXME text decoder
+        logit_scale_exp = self.logit_scale.exp() if output_logits or output_logit_scale_bias else None
+        if output_logit_scale_bias:
+            output["logit_scale"] = logit_scale_exp
+            if self.logit_bias is not None:
+                output['logit_bias'] = self.logit_bias
+
+        return output
+
+    def forward(
+            self,
+            image,
+            text: Optional[torch.Tensor] = None,
+            image_latent: Optional[torch.Tensor] = None,
+            image_embs: Optional[torch.Tensor] = None,
+            output_labels: bool = True,
+    ):
+        if image_latent is None or image_embs is None:
+            image_latent, image_embs = self._encode_image(image)
+
+        if text is None:
+            return {"image_features": image_latent, "image_embs": image_embs}
+
+        text_latent, token_embs = self._encode_text(text)
+
+        # FIXME this isn't an ideal solution, would like to improve -RW
+        labels: Optional[torch.Tensor] = text[:, 1:] if output_labels else None
+        if output_labels:
+            # align text_embs and thus logits with labels for teacher-forcing caption loss
+            token_embs = token_embs[:, :-1]
+
+        logits = self.text_decoder(image_embs, token_embs)
+        out_dict = {
+            "image_features": image_latent,
+            "text_features": text_latent,
+            "logits": logits,
+            "logit_scale": self.logit_scale.exp()
+        }
+        if labels is not None:
+            out_dict["labels"] = labels
+        if self.logit_bias is not None:
+            out_dict["logit_bias"] = self.logit_bias
+        return out_dict
+
+    def generate(
+        self,
+        image,
+        text=None,
+        seq_len=30,
+        max_seq_len=77,
+        temperature=1.,
+        generation_type="beam_search",
+        top_p=0.1,  # keep tokens in the 1 - top_p quantile
+        top_k=1,  # keeps the top_k most probable tokens
+        pad_token_id=None,
+        eos_token_id=None,
+        sot_token_id=None,
+        num_beams=6,
+        num_beam_groups=3,
+        min_seq_len=5,
+        stopping_criteria=None,
+        repetition_penalty=1.0,
+        fixed_output_length=False # if True output.shape == (batch_size, seq_len)
+    ):
+        # taking many ideas and components from HuggingFace GenerationMixin
+        # https://huggingface.co/docs/transformers/main/en/main_classes/text_generation
+        assert _has_transformers, "Please install transformers for generate functionality. `pip install transformers`."
+        assert seq_len > min_seq_len, "seq_len must be larger than min_seq_len"
+        device = image.device
+
+        with torch.no_grad():
+            sot_token_id = _token_to_tensor(49406 if sot_token_id is None else sot_token_id, device=device)
+            eos_token_id = _token_to_tensor(49407 if eos_token_id is None else eos_token_id, device=device)
+            pad_token_id = self.pad_id if pad_token_id is None else pad_token_id
+            logit_processor = LogitsProcessorList(
+                [
+                    MinLengthLogitsProcessor(min_seq_len, eos_token_id),
+                    RepetitionPenaltyLogitsProcessor(repetition_penalty),
+                ]
+            )
+
+            if stopping_criteria is None:
+                stopping_criteria = [MaxLengthCriteria(max_length=seq_len)]
+            stopping_criteria = StoppingCriteriaList(stopping_criteria)
+
+            if generation_type == "beam_search":
+                output = self._generate_beamsearch(
+                    image_inputs=image,
+                    pad_token_id=pad_token_id,
+                    eos_token_id=eos_token_id,
+                    sot_token_id=sot_token_id,
+                    num_beams=num_beams,
+                    num_beam_groups=num_beam_groups,
+                    min_seq_len=min_seq_len,
+                    stopping_criteria=stopping_criteria,
+                    logit_processor=logit_processor,
+                )
+                if fixed_output_length and output.shape[1] < seq_len:
+                    pad_len = seq_len - output.shape[1]
+                    return torch.cat((
+                            output,
+                            torch.ones(output.shape[0], pad_len, device=device, dtype=output.dtype) * pad_token_id
+                        ),
+                        dim=1
+                    )
+                return output
+
+            elif generation_type == "top_p":
+                logit_warper = GENERATION_TYPES[generation_type](top_p)
+            elif generation_type == "top_k":
+                logit_warper = GENERATION_TYPES[generation_type](top_k)
+            else:
+                raise ValueError(
+                    f"generation_type has to be one of "
+                    f"{'| ' + ' | '.join(list(GENERATION_TYPES.keys())) + ' |'}."
+                )
+
+            image_latent, image_embs = self._encode_image(image)
+
+            if text is None:
+                text = torch.ones((image.shape[0], 1), device=device, dtype=torch.long) * sot_token_id
+
+            was_training = self.training
+            num_dims = len(text.shape)
+
+            if num_dims == 1:
+                text = text[None, :]
+
+            self.eval()
+            out = text
+
+            while True:
+                x = out[:, -max_seq_len:]
+                cur_len = x.shape[1]
+                logits = self(
+                    image,
+                    x,
+                    image_latent=image_latent,
+                    image_embs=image_embs,
+                    output_labels=False,
+                )["logits"][:, -1]
+                mask = (out[:, -1] == eos_token_id) | (out[:, -1] == pad_token_id)
+                sample = torch.ones((out.shape[0], 1), device=device, dtype=torch.long) * pad_token_id
+
+                if mask.all():
+                    if not fixed_output_length:
+                        break
+                else:
+                    logits = logits[~mask, :]
+                    filtered_logits = logit_processor(x[~mask, :], logits)
+                    filtered_logits = logit_warper(x[~mask, :], filtered_logits)
+                    probs = F.softmax(filtered_logits / temperature, dim=-1)
+
+                    if (cur_len + 1 == seq_len):
+                        sample[~mask, :] = torch.ones((sum(~mask), 1), device=device, dtype=torch.long) * eos_token_id
+                    else:
+                        sample[~mask, :] = torch.multinomial(probs, 1)
+
+                out = torch.cat((out, sample), dim=-1)
+
+                cur_len += 1
+
+                if all(stopping_criteria(out, None)):
+                    break
+
+            if num_dims == 1:
+                out = out.squeeze(0)
+
+            self.train(was_training)
+            return out
+
+    def _generate_beamsearch(
+            self,
+            image_inputs,
+            pad_token_id=None,
+            eos_token_id=None,
+            sot_token_id=None,
+            num_beams=6,
+            num_beam_groups=3,
+            min_seq_len=5,
+            stopping_criteria=None,
+            logit_processor=None,
+            logit_warper=None,
+    ):
+        device = image_inputs.device
+        batch_size = image_inputs.shape[0]
+        image_inputs = torch.repeat_interleave(image_inputs, num_beams, dim=0)
+        image_latent, image_embs = self._encode_image(image_inputs)
+
+        input_ids = torch.ones((batch_size * num_beams, 1), device=device, dtype=torch.long)
+        input_ids = input_ids * sot_token_id
+        beam_scorer = BeamSearchScorer(
+            batch_size=batch_size,
+            num_beams=num_beams,
+            device=device,
+            num_beam_groups=num_beam_groups,
+        )
+        # instantiate logits processors
+        logits_processor = (
+            LogitsProcessorList([MinLengthLogitsProcessor(min_seq_len, eos_token_id=eos_token_id)])
+            if logit_processor is None
+            else logit_processor
+        )
+
+        num_beams = beam_scorer.num_beams
+        num_beam_groups = beam_scorer.num_beam_groups
+        num_sub_beams = num_beams // num_beam_groups
+        batch_size = len(beam_scorer._beam_hyps) // num_beam_groups
+        batch_beam_size, cur_len = input_ids.shape
+        beam_indices = None
+
+        if num_beams * batch_size != batch_beam_size:
+            raise ValueError(
+                f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}."
+            )
+
+        beam_scores = torch.full((batch_size, num_beams), -1e9, dtype=torch.float, device=device)
+        # initialise score of first beam of each group with 0 and the rest with 1e-9. This ensures that the beams in
+        # the same group don't produce same tokens everytime.
+        beam_scores[:, ::num_sub_beams] = 0
+        beam_scores = beam_scores.view((batch_size * num_beams,))
+
+        while True:
+
+            # predicted tokens in cur_len step
+            current_tokens = torch.zeros(batch_size * num_beams, dtype=input_ids.dtype, device=device)
+
+            # indices which will form the beams in the next time step
+            reordering_indices = torch.zeros(batch_size * num_beams, dtype=torch.long, device=device)
+
+            # do one decoder step on all beams of all sentences in batch
+            model_inputs = prepare_inputs_for_generation(input_ids=input_ids, image_inputs=image_inputs)
+            outputs = self(
+                model_inputs['images'],
+                model_inputs['text'],
+                image_latent=image_latent,
+                image_embs=image_embs,
+                output_labels=False,
+            )
+
+            for beam_group_idx in range(num_beam_groups):
+                group_start_idx = beam_group_idx * num_sub_beams
+                group_end_idx = min(group_start_idx + num_sub_beams, num_beams)
+                group_size = group_end_idx - group_start_idx
+
+                # indices of beams of current group among all sentences in batch
+                batch_group_indices = []
+
+                for batch_idx in range(batch_size):
+                    batch_group_indices.extend(
+                        [batch_idx * num_beams + idx for idx in range(group_start_idx, group_end_idx)]
+                    )
+                group_input_ids = input_ids[batch_group_indices]
+
+                # select outputs of beams of currentg group only
+                next_token_logits = outputs['logits'][batch_group_indices, -1, :]
+                vocab_size = next_token_logits.shape[-1]
+
+                next_token_scores_processed = logits_processor(
+                    group_input_ids, next_token_logits, current_tokens=current_tokens, beam_group_idx=beam_group_idx
+                )
+                next_token_scores = next_token_scores_processed + beam_scores[batch_group_indices].unsqueeze(-1)
+                next_token_scores = next_token_scores.expand_as(next_token_scores_processed)
+
+                # reshape for beam search
+                next_token_scores = next_token_scores.view(batch_size, group_size * vocab_size)
+
+                next_token_scores, next_tokens = torch.topk(
+                    next_token_scores, 2 * group_size, dim=1, largest=True, sorted=True
+                )
+
+                next_indices = torch.div(next_tokens, vocab_size, rounding_mode="floor")
+                next_tokens = next_tokens % vocab_size
+
+                # stateless
+                process_beam_indices = sum(beam_indices, ()) if beam_indices is not None else None
+                beam_outputs = beam_scorer.process(
+                    group_input_ids,
+                    next_token_scores,
+                    next_tokens,
+                    next_indices,
+                    pad_token_id=pad_token_id,
+                    eos_token_id=eos_token_id,
+                    beam_indices=process_beam_indices,
+                    group_index=beam_group_idx,
+                )
+                beam_scores[batch_group_indices] = beam_outputs["next_beam_scores"]
+                beam_next_tokens = beam_outputs["next_beam_tokens"]
+                beam_idx = beam_outputs["next_beam_indices"]
+
+                input_ids[batch_group_indices] = group_input_ids[beam_idx]
+                group_input_ids = torch.cat([group_input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1)
+                current_tokens[batch_group_indices] = group_input_ids[:, -1]
+
+                # (beam_idx // group_size) -> batch_idx
+                # (beam_idx % group_size) -> offset of idx inside the group
+                reordering_indices[batch_group_indices] = (
+                    num_beams * torch.div(beam_idx, group_size, rounding_mode="floor") + group_start_idx + (beam_idx % group_size)
+                )
+
+            input_ids = torch.cat([input_ids, current_tokens.unsqueeze(-1)], dim=-1)
+
+            # increase cur_len
+            cur_len = cur_len + 1
+            if beam_scorer.is_done or all(stopping_criteria(input_ids, None)):
+                break
+
+        final_beam_indices = sum(beam_indices, ()) if beam_indices is not None else None
+        sequence_outputs = beam_scorer.finalize(
+            input_ids,
+            beam_scores,
+            next_tokens,
+            next_indices,
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            max_length=stopping_criteria.max_length,
+            beam_indices=final_beam_indices,
+        )
+        return sequence_outputs['sequences']
+
+
+def prepare_inputs_for_generation(input_ids, image_inputs, past=None, **kwargs):
+    if past:
+        input_ids = input_ids[:, -1].unsqueeze(-1)
+
+    attention_mask = kwargs.get("attention_mask", None)
+    position_ids = kwargs.get("position_ids", None)
+
+    if attention_mask is not None and position_ids is None:
+        # create position_ids on the fly for batch generation
+        position_ids = attention_mask.long().cumsum(-1) - 1
+        position_ids.masked_fill_(attention_mask == 0, 1)
+    else:
+        position_ids = None
+    return {
+        "text": input_ids,
+        "images": image_inputs,
+        "past_key_values": past,
+        "position_ids": position_ids,
+        "attention_mask": attention_mask,
+    }

src/open_clip/factory.py

@@ -0,0 +1,93 @@
+import json
+import logging
+from copy import deepcopy
+from pathlib import Path
+from typing import Optional, Tuple, Union
+
+import torch
+
+from .biosignals_coca_model import BiosignalsCoCa
+from .model import get_cast_dtype, convert_weights_to_lp
+from .tokenizer import SimpleTokenizer, DEFAULT_CONTEXT_LENGTH
+
+_MODEL_CONFIG_PATHS = [Path(__file__).parent / "model_configs/"]
+_MODEL_CONFIGS = {}
+
+
+def _rescan_model_configs():
+    global _MODEL_CONFIGS
+    config_files = []
+    for config_path in _MODEL_CONFIG_PATHS:
+        if config_path.is_dir():
+            config_files.extend(config_path.glob("*.json"))
+    for cf in config_files:
+        with open(cf, "r") as f:
+            model_cfg = json.load(f)
+            if all(a in model_cfg for a in ("embed_dim", "biosignals_cfg", "text_cfg")):
+                _MODEL_CONFIGS[cf.stem] = model_cfg
+
+
+_rescan_model_configs()
+
+
+def get_model_config(model_name: str):
+    return deepcopy(_MODEL_CONFIGS.get(model_name))
+
+
+def create_model(
+    model_name: str,
+    precision: str = "fp32",
+    device: Union[str, torch.device] = "cpu",
+    **model_kwargs,
+) -> BiosignalsCoCa:
+    if isinstance(device, str):
+        device = torch.device(device)
+
+    model_cfg = get_model_config(model_name)
+    if model_cfg is None:
+        raise RuntimeError(f"Model config for '{model_name}' not found. Available: {list(_MODEL_CONFIGS.keys())}")
+
+    model_cfg.pop("custom_text", None)
+    model_cfg.update(model_kwargs)
+
+    cast_dtype = get_cast_dtype(precision)
+    model = BiosignalsCoCa(**model_cfg, cast_dtype=cast_dtype)
+
+    if precision in ("fp16", "bf16"):
+        dtype = torch.float16 if "fp16" in precision else torch.bfloat16
+        model.to(device=device)
+        convert_weights_to_lp(model, dtype=dtype)
+    elif precision in ("pure_fp16", "pure_bf16"):
+        dtype = torch.float16 if "fp16" in precision else torch.bfloat16
+        model.to(device=device, dtype=dtype)
+    else:
+        model.to(device=device)
+
+    model.output_dict = True
+    return model
+
+
+def load_checkpoint(model, checkpoint_path: str, device="cpu"):
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    state_dict = checkpoint.get("state_dict", checkpoint)
+    if next(iter(state_dict)).startswith("module."):
+        state_dict = {k[len("module."):]: v for k, v in state_dict.items()}
+    incompatible = model.load_state_dict(state_dict, strict=False)
+    return incompatible
+
+
+def get_tokenizer(model_name: str = "", context_length: Optional[int] = None, **kwargs):
+    config = get_model_config(model_name) or {}
+    text_cfg = config.get("text_cfg", {})
+    if context_length is None:
+        context_length = text_cfg.get("context_length", DEFAULT_CONTEXT_LENGTH)
+    return SimpleTokenizer(context_length=context_length, **kwargs)
+
+
+def get_input_dtype(precision: str):
+    input_dtype = None
+    if precision in ("bf16", "pure_bf16"):
+        input_dtype = torch.bfloat16
+    elif precision in ("fp16", "pure_fp16"):
+        input_dtype = torch.float16
+    return input_dtype

src/open_clip/model.py

@@ -0,0 +1,943 @@
+""" CLIP Model
+
+Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+import copy
+import logging
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+from functools import partial
+
+from .transformer import (
+    LayerNormFp32,
+    LayerNorm,
+    QuickGELU,
+    Attention,
+    VisionTransformer,
+    TextTransformer,
+    text_global_pool,
+    lock_text_tower,
+    to_2tuple,
+)
+
+
+@dataclass
+class CLIPVisionCfg:
+    layers: Union[Tuple[int, int, int, int], int] = 12
+    width: int = 768
+    head_width: int = 64
+    mlp_ratio: float = 4.0
+    patch_size: int = 16
+    image_size: Union[Tuple[int, int], int] = 224
+
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    patch_dropout: float = 0.  # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
+    attentional_pool: bool = False  # whether to use attentional pooler in the last embedding layer (overrides pool_type)
+    attn_pooler_queries: int = 256  # n_queries for attentional pooler
+    attn_pooler_heads: int = 8  # n heads for attentional_pooling
+    no_ln_pre: bool = False  # disable pre transformer LayerNorm
+    pos_embed_type: str = 'learnable'
+    final_ln_after_pool: bool = False  # apply final LayerNorm after pooling
+    pool_type: str = 'tok'
+    output_tokens: bool = False
+    act_kwargs: Optional[dict] = None
+    norm_kwargs: Optional[dict] = None
+
+    # Custom attention block settings
+    block_type: Optional[str] = None  # attention block type ('default', 'custom'), auto-selects 'custom' if any below features enabled
+    qk_norm: bool = False  # apply layer norm to q and k in attention
+    scaled_cosine_attn: bool = False  # use scaled cosine attention
+    scale_heads: bool = False  # learnable head-specific scale applied to attention logits
+    scale_attn_inner: bool = False  # apply layer norm on attention context, before output projection
+    scale_attn: bool = False  # apply layer norm after full attention block
+    scale_fc: bool = False  # apply layer norm in MLP block
+
+    timm_model_name: Optional[str] = None  # a valid model name overrides layers, width, patch_size
+    timm_model_pretrained: bool = False  # use (imagenet) pretrained weights for named model
+    timm_pool: str = 'avg'  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
+    timm_proj: str = 'linear'  # linear projection for timm model output ('linear', 'mlp', '')
+    timm_proj_bias: bool = False  # enable bias final projection
+    timm_drop: float = 0.  # head dropout
+    timm_drop_path: Optional[float] = None  # backbone stochastic depth
+
+
+@dataclass
+class CLIPTextCfg:
+    context_length: int = 77
+    vocab_size: int = 49408
+    hf_tokenizer_name: Optional[str] = None
+    tokenizer_mode: Optional[str] = None
+    tokenizer_kwargs: Optional[dict] = None
+
+    width: int = 512
+    heads: int = 8
+    layers: int = 12
+    mlp_ratio: float = 4.0
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    embed_cls: bool = False
+    pad_id: int = 0
+    eos_id: int = 2  # only used for when pool_type == 'eos', must match tokenizer eos
+    no_causal_mask: bool = False  # disable causal masking
+    final_ln_after_pool: bool = False  # apply final LayerNorm after pooling
+    pool_type: str = 'argmax'
+    proj_bias: bool = False
+    proj_type: str = 'linear'  # control final text projection, 'none' forces no projection
+    output_tokens: bool = False
+    act_kwargs: dict = None
+    norm_kwargs: dict = None
+
+    # Custom attention block settings
+    block_type: Optional[str] = None  # attention block type ('default', 'custom'), auto-selects 'custom' if any custom features enabled
+    qk_norm: bool = False  # apply layer norm to q and k in attention
+    scaled_cosine_attn: bool = False  # use scaled cosine attention
+    scale_heads: bool = False  # learnable head-specific scale applied to attention logits
+    scale_attn_inner: bool = False  # apply layer norm on attention context, before output projection
+    scale_attn: bool = False  # apply layer norm after full attention block
+    scale_fc: bool = False  # apply layer norm in MLP block
+
+    # HuggingFace specific text tower config
+    hf_model_name: Optional[str] = None
+    hf_model_pretrained: bool = True
+    hf_proj_type: str = 'mlp'
+    hf_pooler_type: str = 'mean_pooler'  # attentional pooling for HF models
+    special_tokens_to_add: Optional[dict] = None  # special tokens to add to tokenizer (e.g., for Pythia)
+
+
+def get_cast_dtype(precision: str):
+    cast_dtype = None
+    if precision == 'bf16':
+        cast_dtype = torch.bfloat16
+    elif precision == 'fp16':
+        cast_dtype = torch.float16
+    return cast_dtype
+
+
+def get_input_dtype(precision: str):
+    input_dtype = None
+    if precision in ('bf16', 'pure_bf16'):
+        input_dtype = torch.bfloat16
+    elif precision in ('fp16', 'pure_fp16'):
+        input_dtype = torch.float16
+    return input_dtype
+
+
+def _build_vision_tower(
+        embed_dim: int,
+        vision_cfg: CLIPVisionCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None
+):
+    if isinstance(vision_cfg, dict):
+        vision_cfg = CLIPVisionCfg(**vision_cfg)
+
+    # OpenAI models are pretrained w/ QuickGELU but native nn.GELU is both faster and more
+    # memory efficient in recent PyTorch releases (>= 1.10).
+    # NOTE: timm models always use native GELU regardless of quick_gelu flag.
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+
+    if vision_cfg.timm_model_name:
+        from .timm_model import TimmModel
+        visual = TimmModel(
+            vision_cfg.timm_model_name,
+            pretrained=vision_cfg.timm_model_pretrained,
+            pool=vision_cfg.timm_pool,
+            proj=vision_cfg.timm_proj,
+            proj_bias=vision_cfg.timm_proj_bias,
+            drop=vision_cfg.timm_drop,
+            drop_path=vision_cfg.timm_drop_path,
+            patch_drop=vision_cfg.patch_dropout if vision_cfg.patch_dropout > 0 else None,
+            embed_dim=embed_dim,
+            image_size=vision_cfg.image_size,
+        )
+    elif isinstance(vision_cfg.layers, (tuple, list)):
+        from .modified_resnet import ModifiedResNet
+        vision_heads = vision_cfg.width * 32 // vision_cfg.head_width
+        visual = ModifiedResNet(
+            layers=vision_cfg.layers,
+            output_dim=embed_dim,
+            heads=vision_heads,
+            image_size=vision_cfg.image_size,
+            width=vision_cfg.width,
+        )
+    else:
+        vision_heads = vision_cfg.width // vision_cfg.head_width
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+        if vision_cfg.norm_kwargs:
+            norm_layer = partial(norm_layer, **vision_cfg.norm_kwargs)
+        if vision_cfg.act_kwargs is not None:
+            act_layer = partial(act_layer, **vision_cfg.act_kwargs)
+
+        visual = VisionTransformer(
+            image_size=vision_cfg.image_size,
+            patch_size=vision_cfg.patch_size,
+            width=vision_cfg.width,
+            layers=vision_cfg.layers,
+            heads=vision_heads,
+            mlp_ratio=vision_cfg.mlp_ratio,
+            ls_init_value=vision_cfg.ls_init_value,
+            patch_dropout=vision_cfg.patch_dropout,
+            attentional_pool=vision_cfg.attentional_pool,
+            attn_pooler_queries=vision_cfg.attn_pooler_queries,
+            attn_pooler_heads=vision_cfg.attn_pooler_heads,
+            pos_embed_type=vision_cfg.pos_embed_type,
+            no_ln_pre=vision_cfg.no_ln_pre,
+            final_ln_after_pool=vision_cfg.final_ln_after_pool,
+            pool_type=vision_cfg.pool_type,
+            output_tokens=vision_cfg.output_tokens,
+            output_dim=embed_dim,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            block_type=vision_cfg.block_type,
+            qk_norm=vision_cfg.qk_norm,
+            scaled_cosine_attn=vision_cfg.scaled_cosine_attn,
+            scale_heads=vision_cfg.scale_heads,
+            scale_attn_inner=vision_cfg.scale_attn_inner,
+            scale_attn=vision_cfg.scale_attn,
+            scale_fc=vision_cfg.scale_fc,
+        )
+
+    return visual
+
+
+
+
+
+def _build_text_tower(
+        embed_dim: int,
+        text_cfg: CLIPTextCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    if isinstance(text_cfg, dict):
+        text_cfg = CLIPTextCfg(**text_cfg)
+
+    if text_cfg.hf_model_name:
+        from .hf_model import HFTextEncoder
+        text = HFTextEncoder(
+            text_cfg.hf_model_name,
+            output_dim=embed_dim,
+            proj_type=text_cfg.hf_proj_type,
+            pooler_type=text_cfg.hf_pooler_type,
+            pretrained=text_cfg.hf_model_pretrained,
+            output_tokens=text_cfg.output_tokens,
+        )
+        
+        # Handle special tokens if configured (e.g., for Pythia)
+        special_tokens_cfg = getattr(text_cfg, 'special_tokens_to_add', None)
+        if special_tokens_cfg:
+            from transformers import AutoTokenizer
+            import logging
+            
+            # Load tokenizer from local cache only (ensures consistency with get_tokenizer())
+            # get_tokenizer() is called first and downloads/caches, we just reuse that exact version
+            tokenizer = AutoTokenizer.from_pretrained(
+                text_cfg.hf_model_name,
+                local_files_only=True
+            )
+            
+            # Store original vocab size before adding new tokens
+            # This is needed to unfreeze new token embeddings after locking
+            original_vocab_size = len(tokenizer)
+            text.original_vocab_size = original_vocab_size
+            
+            tokenizer.add_special_tokens(special_tokens_cfg)
+            
+            # Resize model embeddings to accommodate new tokens
+            # pad_to_multiple_of=64 ensures optimal Tensor Core performance for embedding lookups
+            new_vocab_size = len(tokenizer)
+            text.transformer.resize_token_embeddings(new_vocab_size, pad_to_multiple_of=64)
+            
+            # Store token IDs for use in forward pass
+            if 'additional_special_tokens' in special_tokens_cfg:
+                for token in special_tokens_cfg['additional_special_tokens']:
+                    if token == '<coca_cls>':
+                        text.coca_cls_token_id = tokenizer.convert_tokens_to_ids(token)
+            
+            if 'pad_token' in special_tokens_cfg:
+                text.config.pad_token_id = tokenizer.pad_token_id
+                text.pad_token_id = tokenizer.pad_token_id
+            
+            text.config.vocab_size = new_vocab_size
+            text.vocab_size = new_vocab_size
+            
+            logging.info(f"Added special tokens to {text_cfg.hf_model_name}:")
+            logging.info(f"  Original vocab size: {original_vocab_size}")
+            logging.info(f"  New vocab size: {new_vocab_size}")
+            logging.info(f"  Added {new_vocab_size - original_vocab_size} new tokens")
+            if text.coca_cls_token_id is not None:
+                logging.info(f"  CoCa CLS token ID: {text.coca_cls_token_id}")
+            if text.pad_token_id is not None:
+                logging.info(f"  Pad token ID: {text.pad_token_id}")
+    else:
+        act_layer = QuickGELU if quick_gelu else nn.GELU
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+        if text_cfg.norm_kwargs:
+            norm_layer = partial(norm_layer, **text_cfg.norm_kwargs)
+        if text_cfg.act_kwargs is not None:
+            act_layer = partial(act_layer, **text_cfg.act_kwargs)
+
+        text = TextTransformer(
+            context_length=text_cfg.context_length,
+            vocab_size=text_cfg.vocab_size,
+            width=text_cfg.width,
+            heads=text_cfg.heads,
+            layers=text_cfg.layers,
+            mlp_ratio=text_cfg.mlp_ratio,
+            ls_init_value=text_cfg.ls_init_value,
+            output_dim=embed_dim,
+            embed_cls=text_cfg.embed_cls,
+            no_causal_mask=text_cfg.no_causal_mask,
+            pad_id=text_cfg.pad_id,
+            eos_id=text_cfg.eos_id,
+            pool_type=text_cfg.pool_type,
+            proj_type=text_cfg.proj_type,
+            proj_bias=text_cfg.proj_bias,
+            output_tokens=text_cfg.output_tokens,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            block_type=text_cfg.block_type,
+            qk_norm=text_cfg.qk_norm,
+            scaled_cosine_attn=text_cfg.scaled_cosine_attn,
+            scale_heads=text_cfg.scale_heads,
+            scale_attn_inner=text_cfg.scale_attn_inner,
+            scale_attn=text_cfg.scale_attn,
+            scale_fc=text_cfg.scale_fc,
+        )
+    return text
+
+
+class CLIP(nn.Module):
+    output_dict: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            init_logit_scale: float = np.log(1 / 0.07),
+            init_logit_bias: Optional[float] = None,
+            nonscalar_logit_scale: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            output_dict: bool = False,
+    ):
+        super().__init__()
+        self.output_dict = output_dict
+
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+
+        text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.transformer = text.transformer
+        self.context_length = text.context_length
+        self.vocab_size = text.vocab_size
+        self.token_embedding = text.token_embedding
+        self.positional_embedding = text.positional_embedding
+        self.ln_final = text.ln_final
+        self.text_projection = text.text_projection
+        self.text_pool_type = text.pool_type
+        self.text_eos_id = text.eos_id
+        self.register_buffer('attn_mask', text.attn_mask, persistent=False)
+
+        lshape = [1] if nonscalar_logit_scale else []
+        self.logit_scale = nn.Parameter(torch.ones(lshape) * init_logit_scale)
+        if init_logit_bias is not None:
+            self.logit_bias = nn.Parameter(torch.ones(lshape) * init_logit_bias)
+        else:
+            self.logit_bias = None
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    def lock_text_tower(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True):
+        assert freeze_layer_norm, 'Unfreezing LayerNorm is not supported. LayerNorm treated like other weights.'
+        lock_text_tower(self, unlocked_layers)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.transformer.grad_checkpointing = enable
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        # for timm optimizers, 1d params like logit_scale, logit_bias, ln/bn scale, biases are excluded by default
+        no_wd = {'positional_embedding'}
+        if hasattr(self.visual, 'no_weight_decay'):
+            for n in self.visual.no_weight_decay():
+                no_wd.add('visual.' + n)
+        return no_wd
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        cast_dtype = self.transformer.get_cast_dtype()
+
+        x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.to(cast_dtype)
+        x = self.transformer(x, attn_mask=self.attn_mask)
+        x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
+        x = text_global_pool(x, text, self.text_pool_type, eos_token_id=getattr(self, "text_eos_id", None))
+        if self.text_projection is not None:
+            if isinstance(self.text_projection, nn.Linear):
+                x = self.text_projection(x)
+            else:
+                x = x @ self.text_projection
+
+        return F.normalize(x, dim=-1) if normalize else x
+
+    def get_logits(self, image, text):
+        image_features = self.encode_image(image, normalize=True)
+        text_features = self.encode_text(text, normalize=True)
+        image_logits = self.logit_scale.exp() * image_features @ text_features.T
+        if self.logit_bias is not None:
+            image_logits += self.logit_bias
+        text_logits = image_logits.T
+        return image_logits, text_logits
+
+    def forward_intermediates(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+            image_indices: Optional[Union[int, List[int]]] = None,
+            text_indices: Optional[Union[int, List[int]]] = None,
+            stop_early: bool = False,
+            normalize: bool = True,
+            normalize_intermediates: bool = False,
+            intermediates_only: bool = False,
+            image_output_fmt: str = 'NCHW',
+            image_output_extra_tokens: bool = False,
+            text_output_fmt: str = 'NLC',
+            text_output_extra_tokens: bool = False,
+            output_logits: bool = False,
+            output_logit_scale_bias: bool = False,
+    ) -> Dict[str, Union[torch.Tensor, List[torch.Tensor]]]:
+        """ Forward features that returns intermediates.
+
+        Args:
+            image: Input image tensor
+            text: Input text tensor
+            image_indices: For image tower, Take last n blocks if int, all if None, select matching indices if sequence
+            text_indices: Take last n blocks if int, all if None, select matching indices if sequence
+            stop_early: Stop iterating over blocks when last desired intermediate hit
+            normalize_intermediates: Apply final norm layer to all intermediates
+            normalize: L2 Normalize final features
+            intermediates_only: Only return intermediate features, do not return final features
+            image_output_fmt: Shape of intermediate image feature outputs
+            image_output_extra_tokens: Return both prefix and spatial intermediate tokens
+            text_output_fmt: Shape of intermediate text feature outputs (ignored for this model)
+            text_output_extra_tokens: Return both prefix and spatial intermediate tokens (ignored for this model)
+            output_logits: Include logits in output
+            output_logit_scale_bias: Include the logit scale bias in the output
+        Returns:
+
+        """
+        output = {}
+        if intermediates_only:
+            # intermediates only disables final feature normalization, and include logits
+            normalize = False
+            output_logits = False
+        if output_logits:
+            assert image is not None and text is not None, 'Both image and text inputs are required to compute logits'
+
+        if image is not None:
+            image_output = self.visual.forward_intermediates(
+                image,
+                indices=image_indices,
+                stop_early=stop_early,
+                normalize_intermediates=normalize_intermediates,
+                intermediates_only=intermediates_only,
+                output_fmt=image_output_fmt,
+                output_extra_tokens=image_output_extra_tokens,
+            )
+            if normalize and "image_features" in image_output:
+                image_output["image_features"] = F.normalize(image_output["image_features"], dim=-1)
+            output.update(image_output)
+
+        if text is not None:
+            cast_dtype = self.transformer.get_cast_dtype()
+            x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+            x = x + self.positional_embedding.to(cast_dtype)
+            x, intermediates = self.transformer.forward_intermediates(
+                x,
+                attn_mask=self.attn_mask,
+                indices=text_indices
+            )
+            if normalize_intermediates:
+                intermediates = [self.ln_final(xi) for xi in intermediates]
+
+            # NOTE this model doesn't support cls embed in text transformer, no need for extra intermediate tokens
+            output["text_intermediates"] = intermediates
+
+            if not intermediates_only:
+                x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
+                x = text_global_pool(x, text, self.text_pool_type, eos_token_id=getattr(self, "text_eos_id", None))
+                if self.text_projection is not None:
+                    if isinstance(self.text_projection, nn.Linear):
+                        x = self.text_projection(x)
+                    else:
+                        x = x @ self.text_projection
+                if normalize:
+                    x = F.normalize(x, dim=-1)
+                output["text_features"] = x
+
+        logit_scale_exp = self.logit_scale.exp() if output_logits or output_logit_scale_bias else None
+
+        if output_logits:
+            image_logits = logit_scale_exp * output["image_features"] @ output["text_features"].T
+            if self.logit_bias is not None:
+                image_logits += self.logit_bias
+            text_logits = image_logits.T
+            output["image_logits"] = image_logits
+            output["text_logits"] = text_logits
+
+        if output_logit_scale_bias:
+            output["logit_scale"] = logit_scale_exp
+            if self.logit_bias is not None:
+                output['logit_bias'] = self.logit_bias
+
+        return output
+
+    def forward(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+    ):
+        image_features = self.encode_image(image, normalize=True) if image is not None else None
+        text_features = self.encode_text(text, normalize=True) if text is not None else None
+
+        if self.output_dict:
+            out_dict = {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+            if self.logit_bias is not None:
+                out_dict['logit_bias'] = self.logit_bias
+            return out_dict
+
+        if self.logit_bias is not None:
+            return image_features, text_features, self.logit_scale.exp(), self.logit_bias
+        return image_features, text_features, self.logit_scale.exp()
+
+
+class CustomTextCLIP(nn.Module):
+    output_dict: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            init_logit_scale: float = np.log(1 / 0.07),
+            init_logit_bias: Optional[float] = None,
+            nonscalar_logit_scale: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            output_dict: bool = False,
+    ):
+        super().__init__()
+        self.output_dict = output_dict
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+        self.text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.context_length = self.text.context_length
+        self.vocab_size = self.text.vocab_size
+
+        lshape = [1] if nonscalar_logit_scale else []
+        self.logit_scale = nn.Parameter(torch.ones(lshape) * init_logit_scale)
+        if init_logit_bias is not None:
+            self.logit_bias = nn.Parameter(torch.ones(lshape) * init_logit_bias)
+        else:
+            self.logit_bias = None
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    def lock_text_tower(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True):
+        self.text.lock(unlocked_layers, freeze_layer_norm)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.text.set_grad_checkpointing(enable)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        # for timm optimizers, 1d params like logit_scale, logit_bias, ln/bn scale, biases are excluded by default
+        no_wd = set()
+        if hasattr(self.visual, 'no_weight_decay'):
+            for n in self.visual.no_weight_decay():
+                no_wd.add('visual.' + n)
+        if hasattr(self.text, 'no_weight_decay'):
+            for n in self.text.no_weight_decay():
+                no_wd.add('text.' + n)
+        return no_wd
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        features = self.text(text)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def get_logits(self, image, text):
+        image_features = self.encode_image(image, normalize=True)
+        text_features = self.encode_text(text, normalize=True)
+        image_logits = self.logit_scale.exp() * image_features @ text_features.T
+        if self.logit_bias is not None:
+            image_logits += self.logit_bias
+        text_logits = image_logits.T
+        return image_logits, text_logits
+
+    def forward_intermediates(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+            image_indices: Optional[Union[int, List[int]]] = None,
+            text_indices: Optional[Union[int, List[int]]] = None,
+            stop_early: bool = False,
+            normalize: bool = True,
+            normalize_intermediates: bool = False,
+            intermediates_only: bool = False,
+            image_output_fmt: str = 'NCHW',
+            image_output_extra_tokens: bool = False,
+            text_output_fmt: str = 'NLC',
+            text_output_extra_tokens: bool = False,
+            output_logits: bool = False,
+            output_logit_scale_bias: bool = False,
+    ) -> Dict[str, Union[torch.Tensor, List[torch.Tensor]]]:
+        """ Forward features that returns intermediates.
+
+        Args:
+            image: Input image tensor
+            text: Input text tensor
+            image_indices: For image tower, Take last n blocks if int, all if None, select matching indices if sequence
+            text_indices: Take last n blocks if int, all if None, select matching indices if sequence
+            stop_early: Stop iterating over blocks when last desired intermediate hit
+            normalize: L2 Normalize final image and text features (if present)
+            normalize_intermediates: Apply final encoder norm layer to all intermediates (if possible)
+            intermediates_only: Only return intermediate features, do not return final features
+            image_output_fmt: Shape of intermediate image feature outputs
+            image_output_extra_tokens: Return both prefix and spatial intermediate tokens
+            text_output_fmt: Shape of intermediate text feature outputs
+            text_output_extra_tokens: Return both prefix and spatial intermediate tokens
+            output_logits: Include logits in output
+            output_logit_scale_bias: Include the logit scale bias in the output
+        Returns:
+
+        """
+        output = {}
+        if intermediates_only:
+            # intermediates only disables final feature normalization, and include logits
+            normalize = False
+            output_logits = False
+        if output_logits:
+            assert image is not None and text is not None, 'Both image and text inputs are required to compute logits'
+
+        if image is not None:
+            image_output = self.visual.forward_intermediates(
+                image,
+                indices=image_indices,
+                stop_early=stop_early,
+                normalize_intermediates=normalize_intermediates,
+                intermediates_only=intermediates_only,
+                output_fmt=image_output_fmt,
+                output_extra_tokens=image_output_extra_tokens,
+            )
+            if normalize and "image_features" in image_output:
+                image_output["image_features"] = F.normalize(image_output["image_features"], dim=-1)
+            output.update(image_output)
+
+        if text is not None:
+            text_output = self.text.forward_intermediates(
+                text,
+                indices=text_indices,
+                stop_early=stop_early,
+                normalize_intermediates=normalize_intermediates,
+                intermediates_only=intermediates_only,
+                output_fmt=text_output_fmt,
+                output_extra_tokens=text_output_extra_tokens,
+            )
+            if normalize and "text_features" in text_output:
+                text_output["text_features"] = F.normalize(text_output["text_features"], dim=-1)
+            output.update(text_output)
+
+        logit_scale_exp = self.logit_scale.exp() if output_logits or output_logit_scale_bias else None
+
+        if output_logits:
+            image_logits = logit_scale_exp * output["image_features"] @ output["text_features"].T
+            if self.logit_bias is not None:
+                image_logits += self.logit_bias
+            text_logits = image_logits.T
+            output["image_logits"] = image_logits
+            output["text_logits"] = text_logits
+
+        if output_logit_scale_bias:
+            output["logit_scale"] = logit_scale_exp
+            if self.logit_bias is not None:
+                output['logit_bias'] = self.logit_bias
+
+        return output
+
+    def forward(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+    ):
+        image_features = self.encode_image(image, normalize=True) if image is not None else None
+        text_features = self.encode_text(text, normalize=True) if text is not None else None
+
+        if self.output_dict:
+            out_dict = {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+            if self.logit_bias is not None:
+                out_dict['logit_bias'] = self.logit_bias
+            return out_dict
+
+        if self.logit_bias is not None:
+            return image_features, text_features, self.logit_scale.exp(), self.logit_bias
+        return image_features, text_features, self.logit_scale.exp()
+
+
+def convert_weights_to_lp(model: nn.Module, dtype=torch.float16):
+    """Convert applicable model parameters to low-precision (bf16 or fp16)"""
+
+    def _convert_weights(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.to(dtype)
+            if l.bias is not None:
+                l.bias.data = l.bias.data.to(dtype)
+
+        if isinstance(l, (nn.MultiheadAttention, Attention)):
+            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, None)
+                if tensor is not None:
+                    tensor.data = tensor.data.to(dtype)
+
+        if isinstance(l, (CLIP, TextTransformer)):
+            # convert text nn.Parameter projections
+            attr = getattr(l, "text_projection", None)
+            if attr is not None:
+                attr.data = attr.data.to(dtype)
+
+        if isinstance(l, VisionTransformer):
+            # convert vision nn.Parameter projections
+            attr = getattr(l, "proj", None)
+            if attr is not None:
+                attr.data = attr.data.to(dtype)
+
+    model.apply(_convert_weights)
+
+
+convert_weights_to_fp16 = convert_weights_to_lp  # backwards compat
+
+
+# used to maintain checkpoint compatibility
+def convert_to_custom_text_state_dict(state_dict: dict):
+    if 'text_projection' in state_dict:
+        # old format state_dict, move text tower -> .text
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            if any(k.startswith(p) for p in (
+                'text_projection',
+                'positional_embedding',
+                'token_embedding',
+                'transformer',
+                'ln_final',
+            )):
+                k = 'text.' + k
+            new_state_dict[k] = v
+        return new_state_dict
+    return state_dict
+
+
+def build_model_from_openai_state_dict(
+        state_dict: dict,
+        quick_gelu=True,
+        cast_dtype=torch.float16,
+):
+    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_size = 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_size = 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")))
+
+    vision_cfg = CLIPVisionCfg(
+        layers=vision_layers,
+        width=vision_width,
+        patch_size=vision_patch_size,
+        image_size=image_size,
+    )
+    text_cfg = CLIPTextCfg(
+        context_length=context_length,
+        vocab_size=vocab_size,
+        width=transformer_width,
+        heads=transformer_heads,
+        layers=transformer_layers,
+    )
+    model = CLIP(
+        embed_dim,
+        vision_cfg=vision_cfg,
+        text_cfg=text_cfg,
+        quick_gelu=quick_gelu,  # OpenAI models were trained with QuickGELU
+        cast_dtype=cast_dtype,
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        state_dict.pop(key, None)
+    convert_weights_to_fp16(model)  # OpenAI state dicts are partially converted to float16
+    model.load_state_dict(state_dict)
+    return model.eval()
+
+
+def trace_model(model, batch_size=256, device=torch.device('cpu')):
+    model.eval()
+    image_size = model.visual.image_size
+    example_images = torch.ones((batch_size, 3, image_size, image_size), device=device)
+    example_text = torch.zeros((batch_size, model.context_length), dtype=torch.int, device=device)
+    model = torch.jit.trace_module(
+        model,
+        inputs=dict(
+            forward=(example_images, example_text),
+            encode_text=(example_text,),
+            encode_image=(example_images,)
+        ))
+    model.visual.image_size = image_size
+    return model
+
+
+def resize_pos_embed(state_dict, model, interpolation: str = 'bicubic', antialias: bool = True):
+    # Rescale the grid of position embeddings when loading from state_dict
+    old_pos_embed = state_dict.get('visual.positional_embedding', None)
+    if old_pos_embed is None or not hasattr(model.visual, 'grid_size'):
+        return
+    grid_size = to_2tuple(model.visual.grid_size)
+    extra_tokens = 1  # FIXME detect different token configs (ie no class token, or more)
+    new_seq_len = grid_size[0] * grid_size[1] + extra_tokens
+    if new_seq_len == old_pos_embed.shape[0]:
+        return
+
+    if extra_tokens:
+        pos_emb_tok, pos_emb_img = old_pos_embed[:extra_tokens], old_pos_embed[extra_tokens:]
+    else:
+        pos_emb_tok, pos_emb_img = None, old_pos_embed
+    old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img))))
+
+    logging.info('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size)
+    pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
+    pos_emb_img = F.interpolate(
+        pos_emb_img,
+        size=grid_size,
+        mode=interpolation,
+        antialias=antialias,
+        align_corners=False,
+    )
+    pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)[0]
+    if pos_emb_tok is not None:
+        new_pos_embed = torch.cat([pos_emb_tok, pos_emb_img], dim=0)
+    else:
+        new_pos_embed = pos_emb_img
+    state_dict['visual.positional_embedding'] = new_pos_embed
+
+
+def resize_text_pos_embed(state_dict, model, interpolation: str = 'linear', antialias: bool = False):
+    pos_embed_key = 'positional_embedding' if 'positional_embedding' in state_dict else 'text.positional_embedding'
+    old_pos_embed = state_dict.get(pos_embed_key, None)
+    if old_pos_embed is None:
+        return
+    # FIXME add support for text cls_token
+    model_pos_embed = getattr(model, 'positional_embedding', None)
+    if model_pos_embed is None:
+        model_pos_embed = getattr(model.text, 'positional_embedding', None)
+
+    old_num_pos = old_pos_embed.shape[0]
+    old_width = old_pos_embed.shape[1]
+    num_pos = model_pos_embed.shape[0]
+    width = model_pos_embed.shape[1]
+    assert old_width == width, 'text pos_embed width changed!'
+    if old_num_pos == num_pos:
+        return
+
+    logging.info('Resizing text position embedding num_pos from %s to %s', old_num_pos, num_pos)
+    old_pos_embed = old_pos_embed.reshape(1, old_num_pos, old_width).permute(0, 2, 1)
+    old_pos_embed = F.interpolate(
+        old_pos_embed,
+        size=num_pos,
+        mode=interpolation,
+        antialias=antialias,
+        align_corners=False,
+    )
+    old_pos_embed = old_pos_embed.permute(0, 2, 1)[0]
+    new_pos_embed = old_pos_embed
+
+    state_dict[pos_embed_key] = new_pos_embed
+
+
+def get_model_preprocess_cfg(model):
+    module = getattr(model, 'visual', model)
+    preprocess_cfg = getattr(module, 'preprocess_cfg', {})
+    if not preprocess_cfg:
+        # use separate legacy attributes if preprocess_cfg dict not found
+        size = getattr(module, 'image_size')
+        if size is not None:
+            preprocess_cfg['size'] = size
+        mean = getattr(module, 'image_mean', None)
+        if mean is not None:
+            preprocess_cfg['mean'] = mean
+        std = getattr(module, 'image_std', None)
+        if std is not None:
+            preprocess_cfg['std'] = std
+    return preprocess_cfg
+
+
+def set_model_preprocess_cfg(model, preprocess_cfg: Dict[str, Any]):
+    module = getattr(model, 'visual', model)
+    module.image_mean = preprocess_cfg['mean']  # legacy attribute, keeping for bwd compat
+    module.image_std = preprocess_cfg['std']  # legacy attribute, keeping for bwd compat
+    module.preprocess_cfg = copy.deepcopy(preprocess_cfg)  # new attr, package all pp cfg as dict
+
+
+def get_model_tokenize_cfg(model):
+    module = getattr(model, 'text', model)
+    cfg = {}
+    context_length = getattr(module, 'context_length', None)
+    if context_length is not None:
+        cfg['context_length'] = context_length
+    vocab_size = getattr(module, 'vocab_size', None)
+    if vocab_size is not None:
+        cfg['vocab_size'] = vocab_size
+    return cfg

src/open_clip/model_configs/sleep_coca_base_dualtransformer.json

@@ -0,0 +1,44 @@
+{
+    "embed_dim": 512,
+    "multimodal_cfg": {
+        "width": 768,
+        "context_length": 256,
+        "mlp_ratio": 4,
+        "layers": 12,
+        "heads": 12
+    },
+    "biosignals_cfg": {
+        "architecture": "pure_transformer",
+        "input_channels": 10,
+        "signal_length": 1920,
+        "sampling_rate": 64,
+        "patch_size": 16,
+        "conv_embed_dim": 256,
+        "num_temporal_layers": 1,
+        "activation": "swiglu",
+        "norm_type": "rmsnorm",
+        "mlp_bias": false,
+        "share_channel_rope": true,
+        "transformer_layers": 3,
+        "transformer_width": 768,
+        "transformer_heads": 12,
+        "mlp_ratio": 3.0,
+        "pool_type": "attn",
+        "dropout": 0.1,
+        "decoder_tokens": 32
+    },
+    "text_cfg": {
+        "context_length": 256,
+        "vocab_size": 49408,
+        "layers": 12,
+        "heads": 12, 
+        "width": 768,
+        "embed_cls": true,
+        "output_tokens": true
+    },
+    "custom_text": true,
+    "prefix_len": 1,
+    "num_caption_channels": 12,
+    "decoder_type": "cross_attention"
+}
+

src/open_clip/tokenizer.py

@@ -0,0 +1,621 @@
+""" CLIP tokenizer
+
+Copied from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+import gzip
+import html
+import os
+import random
+import string
+from functools import lru_cache, partial
+from typing import Callable, List, Optional, Union, Dict
+import warnings
+
+import ftfy
+import numpy as np
+import regex as re
+import torch
+
+# https://stackoverflow.com/q/62691279
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+_nltk_init = False
+
+DEFAULT_CONTEXT_LENGTH = 77  # default context length for OpenAI CLIP
+
+
+@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 significant 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 = " ".join(text.split())
+    text = text.strip()
+    return text
+
+
+def _clean_canonicalize(x):
+    # basic, remove whitespace, remove punctuation, lower case
+    return canonicalize_text(basic_clean(x))
+
+
+def _clean_lower(x):
+    # basic, remove whitespace, lower case
+    return whitespace_clean(basic_clean(x)).lower()
+
+
+def _clean_whitespace(x):
+    # basic, remove whitespace
+    return whitespace_clean(basic_clean(x))
+
+
+def get_clean_fn(type: str):
+    if type == 'canonicalize':
+        return _clean_canonicalize
+    elif type == 'lower':
+        return _clean_lower
+    elif type == 'whitespace':
+        return _clean_whitespace
+    else:
+        assert False, f"Invalid clean function ({type})."
+
+
+def canonicalize_text(
+    text,
+    *,
+    keep_punctuation_exact_string=None,
+    trans_punctuation: dict = str.maketrans("", "", string.punctuation),
+):
+    """Returns canonicalized `text` (lowercase and punctuation removed).
+
+    From: https://github.com/google-research/big_vision/blob/53f18caf27a9419231bbf08d3388b07671616d3d/big_vision/evaluators/proj/image_text/prompt_engineering.py#L94
+
+    Args:
+      text: string to be canonicalized.
+      keep_punctuation_exact_string: If provided, then this exact string kept.
+        For example providing '{}' will keep any occurrences of '{}' (but will
+        still remove '{' and '}' that appear separately).
+    """
+    text = text.replace("_", " ")
+    if keep_punctuation_exact_string:
+        text = keep_punctuation_exact_string.join(
+            part.translate(trans_punctuation)
+            for part in text.split(keep_punctuation_exact_string)
+        )
+    else:
+        text = text.translate(trans_punctuation)
+    text = text.lower()
+    text = " ".join(text.split())
+    return text.strip()
+
+
+class SimpleTokenizer(object):
+    def __init__(
+            self,
+            bpe_path: str = default_bpe(),
+            additional_special_tokens: Optional[List[str]] = None,
+            context_length: Optional[int] = DEFAULT_CONTEXT_LENGTH,
+            clean: str = 'lower',
+            reduction_mask: str = ''
+    ):
+        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))
+        special_tokens = ['<start_of_text>', '<end_of_text>']
+        if additional_special_tokens:
+            special_tokens += additional_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]
+        self.sot_token_id = self.all_special_ids[0]
+        self.eot_token_id = self.all_special_ids[1]
+        self.context_length = context_length
+        self.clean_fn = get_clean_fn(clean)
+        self.reduction_fn = get_reduction_mask_fn(reduction_mask) if reduction_mask else None
+
+    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 Exception:
+                    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 = self.clean_fn(text)
+        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
+
+    def __call__(self, texts: Union[str, List[str]], context_length: Optional[int] = None) -> 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]
+
+        context_length = context_length or self.context_length
+        assert context_length, 'Please set a valid context length'
+
+        if self.reduction_fn is not None:
+            # use reduction strategy for tokenize if set, otherwise default to truncation below
+            return self.reduction_fn(
+                texts,
+                context_length=context_length,
+                sot_token_id=self.sot_token_id,
+                eot_token_id=self.eot_token_id,
+                encode_fn=self.encode,
+            )
+
+        all_tokens = [[self.sot_token_id] + self.encode(text) + [self.eot_token_id] 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:
+                tokens = tokens[:context_length]  # Truncate
+                tokens[-1] = self.eot_token_id
+            result[i, :len(tokens)] = torch.tensor(tokens)
+
+        return result
+
+
+_tokenizer = SimpleTokenizer()
+
+
+def decode(output_ids: torch.Tensor):
+    output_ids = output_ids.cpu().numpy()
+    return _tokenizer.decode(output_ids)
+
+
+def tokenize(texts: Union[str, List[str]], context_length: int = DEFAULT_CONTEXT_LENGTH) -> torch.LongTensor:
+    return _tokenizer(texts, context_length=context_length)
+
+
+def random_mask_tokenize(
+        texts: Union[str, List[str]],
+        context_length: int,
+        sot_token_id: int,
+        eot_token_id: int,
+        encode_fn: Callable,
+        shuffle: bool = False,
+):
+    all_tokens = [encode_fn(text) for text in texts]
+    result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+    for i, tokens in enumerate(all_tokens):
+        tokens = torch.tensor(tokens)
+        num_tokens = len(tokens)
+        if num_tokens > context_length - 2:  # 2 for sot and eot token
+            num_keep = context_length - 2
+            indices = torch.randperm(len(tokens))
+            indices = indices[:num_keep]
+            if not shuffle:
+                indices = indices.msort()
+            tokens = tokens[indices]
+            num_tokens = num_keep
+        result[i, 0] = sot_token_id
+        result[i, 1:num_tokens + 1] = tokens
+        result[i, num_tokens + 1] = eot_token_id
+
+    return result
+
+
+def simple_mask_tokenize(
+        texts: Union[str, List[str]],
+        context_length: int,
+        sot_token_id: int,
+        eot_token_id: int,
+        encode_fn: Callable,
+):
+    all_tokens = [encode_fn(text) for text in texts]
+    result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+    for i, tokens in enumerate(all_tokens):
+        num_tokens = len(tokens)
+        if num_tokens > context_length - 2:  # 2 for sot and eot token
+            num_keep = context_length - 2
+            start_index = random.randint(0, num_tokens - num_keep)  # high is incl
+            tokens = tokens[start_index: start_index + num_keep]
+        tokens = [sot_token_id] + tokens + [eot_token_id]
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result
+
+
+def syntax_mask_tokenize(
+        texts: Union[str, List[str]],
+        context_length: int,
+        sot_token_id: int,
+        eot_token_id: int,
+        encode_fn: Callable,
+) -> torch.LongTensor:
+    """ Returns the tokenized representation of given input string(s).
+    Apply syntax masking before tokenize.
+    """
+    import nltk
+    global _nltk_init
+    if not _nltk_init:
+        # run them for the first time
+        nltk.download('punkt')
+        nltk.download('averaged_perceptron_tagger')
+        _nltk_init = True
+
+    def get_order(x):
+        if x.startswith('NN'):
+            return 1
+        elif x.startswith('JJ'):
+            return 2
+        elif x.startswith('VB'):
+            return 3
+        else:
+            return 4
+
+    # syntax masking
+    new_texts = []
+    for text in texts:
+        list_tokens = nltk.tokenize.word_tokenize(text)
+        pos_tags = nltk.pos_tag(list_tokens)
+        #  sample the words by get_order method
+        order_list = [get_order(tag) for _, tag in pos_tags]
+        sorted_ids = np.argsort(np.array(order_list))
+        sampled_ids = sorted(sorted_ids[:context_length - 2]) # need 2 slots for sot and eot tokens
+        sampled_tokens = np.take(np.array(list_tokens), sampled_ids, axis=0)  # sample the tokens
+
+        new_text = ''
+        for token in sampled_tokens:
+            new_text = new_text + str(token) + ' '
+        new_text = new_text.strip()
+        new_texts.append(new_text)
+    texts = new_texts
+
+    all_tokens = [[sot_token_id] + encode_fn(text) + [eot_token_id] for text in texts]
+    result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+    for i, tokens in enumerate(all_tokens):
+        # still need first truncate because some words produces two tokens
+        if len(tokens) > context_length:
+            tokens = tokens[:context_length]  # Truncate
+            tokens[-1] = eot_token_id
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result
+
+
+def get_reduction_mask_fn(type: str):
+    """ Choose strategy for dropping (masking) tokens to achieve target context length"""
+    assert type in ('simple', 'random', 'shuffle', 'syntax')
+    if type == 'simple':
+        return simple_mask_tokenize  # randomly select block [start:end]
+    elif type == 'random':
+        return random_mask_tokenize  # randomly drop tokens (keep order)
+    elif type == 'shuffle':
+        return partial(random_mask_tokenize, shuffle=True)  # randomly drop tokens (shuffle order)
+    elif type == 'syntax':
+        return syntax_mask_tokenize  # randomly drop prioritized by syntax
+    else:
+        assert False, F'Unknown type {type}.'
+
+
+class HFTokenizer:
+    """HuggingFace tokenizer wrapper with support for custom tokenization modes"""
+
+    def __init__(
+            self,
+            tokenizer_name: str,
+            context_length: Optional[int] = DEFAULT_CONTEXT_LENGTH,
+            clean: str = 'whitespace',
+            strip_sep_token: bool = False,
+            language: Optional[str] = None,
+            cache_dir: Optional[str] = None,
+            tokenizer_mode: Optional[str] = None,  # None, 'clips'
+            **kwargs
+    ):
+        self.tokenizer_mode = tokenizer_mode or ''
+        self.context_length = context_length
+        self.clean_fn = get_clean_fn(clean)
+        self.strip_sep_token = strip_sep_token
+
+        # NOTE: Left as example of loading custom tokenizer from file for experimentation
+        # if self.tokenizer_mode == 'bert_clips':
+        #     self.special_tokens = {
+        #         "bos_token": 1,
+        #         "eos_token": 2,
+        #         "cls_token": 101,
+        #         "pad_token": 0
+        #     }
+        #
+        #     # For BERT CLIPS mode with vocab file
+        #     from tokenizers import BertWordPieceTokenizer
+        #     if tokenizer_name.startswith('hf-hub:'):
+        #         from huggingface_hub import hf_hub_download
+        #         # Format: hf-hub:repo_id/filename
+        #         repo_url = tokenizer_name[7:]
+        #         parts = repo_url.split('/')
+        #         filename = parts[-1]
+        #         repo_id = '/'.join(parts[:-1])
+        #         vocab_file = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=cache_dir)
+        #         self.tokenizer = BertWordPieceTokenizer(lowercase=True)
+        #         self.tokenizer = self.tokenizer.from_file(vocab_file)
+        #     else:
+        #         # Assume tokenizer_name is a local path to a vocab file
+        #         self.tokenizer = BertWordPieceTokenizer(lowercase=True)
+        #         self.tokenizer = self.tokenizer.from_file(tokenizer_name)
+
+        # Standard HuggingFace tokenizer initialization
+        from transformers import AutoTokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_name,
+            cache_dir=cache_dir,
+            **kwargs
+        )
+
+        # Set language function if available
+        set_lang_fn = getattr(self.tokenizer, 'set_src_lang_special_tokens', None)
+        if callable(set_lang_fn):
+            self.set_lang_fn = set_lang_fn
+        if language is not None:
+            self.set_language(language)
+
+    def save_pretrained(self, dest):
+        self.tokenizer.save_pretrained(dest)
+
+    def __call__(self, texts: Union[str, List[str]], context_length: Optional[int] = None) -> torch.Tensor:
+        # same cleaning as for default tokenizer, except lowercasing
+        # adding lower (for case-sensitive tokenizers) will make it more robust but less sensitive to nuance
+        if isinstance(texts, str):
+            texts = [texts]
+
+        context_length = context_length or self.context_length
+        assert context_length, 'Please set a valid context length in class init or call.'
+
+        texts = [self.clean_fn(text) for text in texts]
+
+        # Handle different tokenization modes
+        if self.tokenizer_mode == 'clips':
+            return self._clips_tokenize(texts, context_length)
+        else:
+            # Standard tokenization
+            input_ids = self.tokenizer.batch_encode_plus(
+                texts,
+                return_tensors='pt',
+                max_length=context_length,
+                padding='max_length',
+                truncation=True,
+            ).input_ids
+
+            if self.strip_sep_token:
+                input_ids = torch.where(
+                    input_ids == self.tokenizer.sep_token_id,
+                    torch.zeros_like(input_ids),
+                    input_ids,
+                )
+
+            return input_ids
+
+    def set_language(self, src_lang):
+        if hasattr(self, 'set_lang_fn'):
+            self.set_lang_fn(src_lang)
+        else:
+            warnings.warn('Cannot set language for the tokenizer.')
+
+    def _clips_tokenize(self, texts: List[str], context_length: int) -> torch.Tensor:
+        """Use standard HF tokenizer but apply custom post-processing"""
+        # Use standard tokenizer without special tokens - we'll add our own
+        encoded_outputs = self.tokenizer.batch_encode_plus(
+            texts,
+            add_special_tokens=False,
+            padding=False,
+            truncation=False,
+            return_tensors=None
+        )
+
+        encoded = []
+        for tokens in encoded_outputs["input_ids"]:
+            tokens = tokens[:context_length - 3]  # Leave room for special tokens
+            tokens = [self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id]
+            encoded.append(tokens)
+
+        # Create result tensor and handle padding + class token
+        result = torch.zeros(len(encoded), context_length, dtype=torch.long)
+        for i, tokens in enumerate(encoded):
+            padded_tokens = self._pad_and_add_class_token(
+                tokens,
+                max_length=context_length,
+                pad_token_id=self.tokenizer.pad_token_id,
+                cls_token_id=self.tokenizer.cls_token_id,
+            )
+            result[i, :len(padded_tokens)] = torch.tensor(padded_tokens)
+
+        return result
+
+    def _pad_and_add_class_token(
+            self,
+            tokens: List[int],
+            max_length: int,
+            pad_token_id: int = 0,
+            cls_token_id: int = 101,
+    ) -> List[int]:
+        """ Add padding with class token at the end """
+        if len(tokens) > max_length - 1:
+            tokens = tokens[:max_length - 1]
+
+        # Add padding to reach max_length-1
+        if len(tokens) < max_length - 1:
+            tokens = tokens + [pad_token_id] * (max_length - 1 - len(tokens))
+
+        # Add class token at the end
+        tokens = tokens + [cls_token_id]
+        return tokens
+
+
+class SigLipTokenizer:
+    """HuggingFace tokenizer wrapper for SigLIP T5 compatible sentencepiece vocabs
+
+    NOTE: this is not needed in normal library use, but is used to import new sentencepiece tokenizers
+    into OpenCLIP. Leaving code here in case future models use new tokenizers.
+    """
+    VOCAB_FILES = {
+        # english, vocab_size=32_000
+        "c4-en": "http://storage.googleapis.com/t5-data/vocabs/cc_en.32000/sentencepiece.model",
+        # used in multilingual models (mT5, PaLI), vocab_size=250_000
+        "mc4": "http://storage.googleapis.com/t5-data/vocabs/mc4.250000.100extra/sentencepiece.model",
+        # used in SigLIP2 models, vocab_size=256000
+        "gemma": "http://storage.googleapis.com/big_vision/gemma_tokenizer.model",
+    }
+
+    def __init__(
+            self,
+            tokenizer_name: str,
+            context_length: Optional[int] = 64,
+    ):
+        if 'gemma' in tokenizer_name:
+            from transformers import GemmaTokenizerFast
+            tokenizer_cls = partial(
+                GemmaTokenizerFast, padding_side='right', add_bos_token=False, add_eos_token=True)
+        else:
+            from transformers import T5TokenizerFast
+            tokenizer_cls = partial(T5TokenizerFast, extra_ids=0)
+
+        if tokenizer_name in self.VOCAB_FILES:
+            # FIXME temporary hack?
+            import tempfile
+            import fsspec
+            vocab_file = self.VOCAB_FILES[tokenizer_name]
+            with tempfile.NamedTemporaryFile('wb') as dst:
+                with fsspec.open(vocab_file, 'rb') as src:
+                    dst.write(src.read())
+                self.tokenizer = tokenizer_cls(dst.name, legacy=False)
+        else:
+            self.tokenizer = tokenizer_cls(tokenizer_name, legacy=False)
+
+        self.tokenizer.pad_token_id = 0 if 'gemma' in tokenizer_name else 1
+        self.tokenizer.eos_token_id = 1
+        self.context_length = context_length
+
+    def save_pretrained(self, dest):
+        self.tokenizer.save_pretrained(dest)
+
+    def __call__(self, texts: Union[str, List[str]], context_length: Optional[int] = None) -> torch.Tensor:
+        # same cleaning as for default tokenizer, except lowercasing
+        # adding lower (for case-sensitive tokenizers) will make it more robust but less sensitive to nuance
+        if isinstance(texts, str):
+            texts = [texts]
+
+        context_length = context_length or self.context_length
+        assert context_length, 'Please set a valid context length in class init or call.'
+
+        texts = [canonicalize_text(basic_clean(text)) for text in texts]
+        output = self.tokenizer(
+            texts,
+            return_tensors='pt',
+            max_length=context_length,
+            padding='max_length',
+            truncation=True,
+        )
+        return output.input_ids

src/open_clip/transformer.py

@@ -0,0 +1,1823 @@
+from collections import OrderedDict
+import math
+from typing import Callable, Dict, List, Optional, Sequence, Tuple, Type, Union
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint
+
+import warnings
+import numpy as np
+
+
+def to_2tuple(x):
+    if isinstance(x, (tuple, list)):
+        return x
+    return (x, x)
+
+
+def feature_take_indices(num_blocks, indices):
+    take_indices = [i if i >= 0 else num_blocks + i for i in indices]
+    max_index = max(take_indices)
+    return take_indices, max_index
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)
+    grid = np.stack(grid, axis=0).reshape([2, 1, grid_size, grid_size])
+    pos_embed = _get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def _get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+    emb_h = _get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
+    emb_w = _get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
+    return np.concatenate([emb_h, emb_w], axis=1)
+
+
+def _get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega
+    pos = pos.reshape(-1)
+    out = np.einsum('m,d->md', pos, omega)
+    return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+
+
+class LayerNormFp32(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16 (by casting to float32 and back)."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        x = F.layer_norm(x.to(torch.float32), self.normalized_shape, self.weight, self.bias, self.eps)
+        return x.to(orig_type)
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm (with cast back to input dtype)."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        return x.to(orig_type)
+
+
+class QuickGELU(nn.Module):
+    # NOTE This is slower than nn.GELU or nn.SiLU and uses more GPU memory
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class LayerScale(nn.Module):
+    def __init__(self, dim, init_values=1e-5, inplace=False):
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x):
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+class PatchDropout(nn.Module):
+    """
+    https://arxiv.org/abs/2212.00794
+    """
+
+    def __init__(
+            self,
+            prob: float = 0.5,
+            exclude_first_token: bool = True
+    ):
+        super().__init__()
+        assert 0 <= prob < 1.
+        self.prob = prob
+        self.exclude_first_token = exclude_first_token  # exclude CLS token
+
+    def forward(self, x):
+        if not self.training or self.prob == 0.:
+            return x
+
+        if self.exclude_first_token:
+            cls_tokens, x = x[:, :1], x[:, 1:]
+        else:
+            cls_tokens = torch.jit.annotate(torch.Tensor, x[:, :1])
+
+        batch = x.size()[0]
+        num_tokens = x.size()[1]
+
+        batch_indices = torch.arange(batch)
+        batch_indices = batch_indices[..., None]
+
+        keep_prob = 1 - self.prob
+        num_patches_keep = max(1, int(num_tokens * keep_prob))
+
+        rand = torch.randn(batch, num_tokens)
+        patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices
+
+        x = x[batch_indices, patch_indices_keep]
+
+        if self.exclude_first_token:
+            x = torch.cat((cls_tokens, x), dim=1)
+
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int = 8,
+            qkv_bias: bool = True,
+            qk_norm: bool = False,
+            scaled_cosine: bool = False,
+            scale_heads: bool = False,
+            inner_norm: bool = False,
+            logit_scale_max: float = math.log(1. / 0.01),
+            norm_layer: Type[nn.Module] = LayerNormFp32,
+            attn_drop: float = 0.,
+            proj_drop: float = 0.
+    ):
+        super().__init__()
+        assert not (scaled_cosine and qk_norm), "Cannot activate both scaled cosine and QK normalization"
+        self.scaled_cosine = scaled_cosine
+        self.scale_heads = scale_heads
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+        self.logit_scale_max = logit_scale_max
+        self.use_fsdpa = hasattr(nn.functional, 'scaled_dot_product_attention')
+
+        # keeping in_proj in this form (instead of nn.Linear) to match weight scheme of original
+        self.in_proj_weight = nn.Parameter(torch.randn((dim * 3, dim)) * self.scale)
+        if qkv_bias:
+            self.in_proj_bias = nn.Parameter(torch.zeros(dim * 3))
+        else:
+            self.in_proj_bias = None
+
+        # QK normalization (with LN) from https://arxiv.org/abs/2106.04560 and related to other QK Norm ideas
+        if qk_norm:
+            self.ln_q = norm_layer(self.head_dim)
+            self.ln_k = norm_layer(self.head_dim)
+        else:
+            self.ln_q = nn.Identity()
+            self.ln_k = nn.Identity()
+
+        # Scaled cosine attention (from Swin Transformer V2, https://arxiv.org/abs/2111.09883)
+        if self.scaled_cosine:
+            self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))))
+        else:
+            self.logit_scale = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+
+        # Per-head attention logit scaling (from NormFormer, https://arxiv.org/abs/2110.09456)
+        if self.scale_heads:
+            self.head_scale = nn.Parameter(torch.ones((num_heads, 1, 1)))
+        else:
+            self.head_scale = None
+
+        # Normalization of attention logits, before final projection.
+        # Origin likely Sub-LN in (Foundation Transformers, https://arxiv.org/abs/2210.06423)
+        if inner_norm:
+            self.ln_inner = norm_layer(dim)
+        else:
+            self.ln_inner = nn.Identity()
+
+        self.out_proj = nn.Linear(dim, dim)
+        self.out_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, attn_mask: Optional[torch.Tensor] = None):
+        N, L, C = x.shape
+        q, k, v = F.linear(x, self.in_proj_weight, self.in_proj_bias).chunk(3, dim=-1)
+        q = q.reshape(N, L, self.num_heads, -1).transpose(1, 2)
+        k = k.reshape(N, L, self.num_heads, -1).transpose(1, 2)
+        v = v.reshape(N, L, self.num_heads, -1).transpose(1, 2)
+
+        if attn_mask is not None:
+            if attn_mask.ndim == 3:
+                # this module works with (L, L), or (N, num_heads, L, L) masks
+                attn_mask = attn_mask.reshape(N, self.num_heads, L, L)
+            if attn_mask.dtype == torch.bool:
+                new_attn_mask = torch.zeros_like(attn_mask, dtype=q.dtype)
+                new_attn_mask.masked_fill_(attn_mask, float("-inf"))
+                attn_mask = new_attn_mask
+            else:
+                attn_mask = attn_mask.to(dtype=q.dtype)
+
+        if self.logit_scale is not None:
+            attn = torch.bmm(
+                F.normalize(q, dim=-1),
+                F.normalize(k, dim=-1).transpose(-1, -2)
+            )
+            logit_scale = torch.clamp(self.logit_scale, max=self.logit_scale_max).exp()
+            attn = attn * logit_scale
+            if attn_mask is not None:
+                attn = attn + attn_mask
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            x = torch.bmm(attn, v)
+        else:
+            q = self.ln_q(q)
+            k = self.ln_k(k)
+            if self.use_fsdpa:
+                x = F.scaled_dot_product_attention(
+                    q, k, v,
+                    attn_mask=attn_mask,
+                    dropout_p=self.attn_drop.p if self.training else 0.,
+                )
+            else:
+                q = q * self.scale
+                attn = torch.bmm(q, k.transpose(-1, -2))
+                if attn_mask is not None:
+                    attn += attn_mask
+                attn = attn.softmax(dim=-1)
+                attn = self.attn_drop(attn)
+                x = torch.bmm(attn, v)
+
+        # N, num_heads, L, head_dim
+        if self.head_scale is not None:
+            x = x * self.head_scale
+        x = x.transpose(1, 2).reshape(N, L, C)
+        x = self.ln_inner(x)
+        x = self.out_proj(x)
+        x = self.out_drop(x)
+        return x
+
+
+class AttentionalPooler(nn.Module):
+    def __init__(
+            self,
+            d_model: int,
+            context_dim: int,
+            n_head: int = 8,
+            n_queries: int = 256,
+            norm_layer: Callable = LayerNorm,
+    ):
+        super().__init__()
+        self.query = nn.Parameter(torch.randn(n_queries, d_model))
+        self.attn = nn.MultiheadAttention(d_model, n_head, kdim=context_dim, vdim=context_dim, batch_first=True)
+        self.ln_q = norm_layer(d_model)
+        self.ln_k = norm_layer(context_dim)
+
+    def forward(self, x: torch.Tensor):
+        N = x.shape[0]
+        x = self.ln_k(x)
+        q = self.ln_q(self.query)
+        out = self.attn(q.unsqueeze(0).expand(N, -1, -1), x, x, need_weights=False)[0]
+        return out
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(
+            self,
+            d_model: int,
+            n_head: int,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            is_cross_attention: bool = False,
+            batch_first: bool = True,
+    ):
+        super().__init__()
+
+        self.ln_1 = norm_layer(d_model)
+        self.attn = nn.MultiheadAttention(d_model, n_head, batch_first=batch_first)
+        self.ls_1 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+        if is_cross_attention:
+            self.ln_1_kv = norm_layer(d_model)
+
+        self.ln_2 = norm_layer(d_model)
+        mlp_width = int(d_model * mlp_ratio)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, mlp_width)),
+            ("gelu", act_layer()),
+            ("c_proj", nn.Linear(mlp_width, d_model))
+        ]))
+        self.ls_2 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+
+    def get_weight_dtype(self) -> torch.dtype:
+        if hasattr(self.mlp.c_fc, 'int8_original_dtype'):
+            return self.mlp.c_fc.int8_original_dtype
+        return self.mlp.c_fc.weight.dtype
+
+    def attention(
+            self,
+            q_x: torch.Tensor,
+            k_x: Optional[torch.Tensor] = None,
+            v_x: Optional[torch.Tensor] = None,
+            attn_mask: Optional[torch.Tensor] = None,
+    ):
+        k_x = k_x if k_x is not None else q_x
+        v_x = v_x if v_x is not None else q_x
+
+        attn_mask = attn_mask.to(q_x.dtype) if attn_mask is not None else None
+        return self.attn(
+            q_x, k_x, v_x,
+            need_weights=False,
+            attn_mask=attn_mask
+        )[0]
+
+    def forward(
+            self,
+            q_x: torch.Tensor,
+            k_x: Optional[torch.Tensor] = None,
+            v_x: Optional[torch.Tensor] = None,
+            attn_mask: Optional[torch.Tensor] = None,
+    ):
+        k_x = self.ln_1_kv(k_x) if hasattr(self, "ln_1_kv") and k_x is not None else None
+        v_x = self.ln_1_kv(v_x) if hasattr(self, "ln_1_kv") and v_x is not None else None
+        x = q_x + self.ls_1(self.attention(q_x=self.ln_1(q_x), k_x=k_x, v_x=v_x, attn_mask=attn_mask))
+        x = x + self.ls_2(self.mlp(self.ln_2(x)))
+        return x
+
+
+class CustomResidualAttentionBlock(nn.Module):
+    def __init__(
+            self,
+            d_model: int,
+            n_head: int,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Type[nn.Module] = nn.GELU,
+            norm_layer: Type[nn.Module] = LayerNorm,
+            qk_norm: bool = False,
+            scale_cosine_attn: bool = False,
+            scale_heads: bool = False,
+            scale_attn_inner: bool = False,
+            scale_attn: bool = False,
+            scale_fc: bool = False,
+            batch_first: bool = True,
+    ):
+        super().__init__()
+        assert batch_first, 'batch_first must be True for CustomResidualAttentionBlock'
+
+        self.ln_1 = norm_layer(d_model)
+        self.attn = Attention(
+            d_model,
+            n_head,
+            qk_norm=qk_norm,
+            scaled_cosine=scale_cosine_attn,
+            scale_heads=scale_heads,
+            inner_norm=scale_attn_inner,
+            norm_layer=norm_layer,
+        )
+        self.ln_attn = norm_layer(d_model) if scale_attn else nn.Identity()
+        self.ls_1 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+
+        self.ln_2 = norm_layer(d_model)
+        mlp_width = int(d_model * mlp_ratio)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, mlp_width)),
+            ("gelu", act_layer()),
+            ('ln', norm_layer(mlp_width) if scale_fc else nn.Identity()),  # from NormFormer / Foundation Transformers
+            ("c_proj", nn.Linear(mlp_width, d_model))
+        ]))
+        self.ls_2 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+
+    def get_weight_dtype(self) -> torch.dtype:
+        if hasattr(self.mlp.c_fc, 'int8_original_dtype'):
+            return self.mlp.c_fc.int8_original_dtype
+        return self.mlp.c_fc.weight.dtype
+
+    def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
+        x = x + self.ls_1(self.ln_attn(self.attn(self.ln_1(x), attn_mask=attn_mask)))
+        x = x + self.ls_2(self.mlp(self.ln_2(x)))
+        return x
+
+
+class CustomTransformer(nn.Module):
+    """ A custom transformer that can use different block types. """
+    def __init__(
+            self,
+            width: int,
+            layers: int,
+            heads: int,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Type[nn.Module] = nn.GELU,
+            norm_layer: Type[nn.Module] = LayerNorm,
+            batch_first: bool = True,
+            block_types: Union[str, List[str]] = 'CustomResidualAttentionBlock',
+    ):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.batch_first = batch_first  # run transformer stack in batch first (N, L, D)
+        self.grad_checkpointing = False
+
+        if isinstance(block_types, str):
+            block_types = [block_types] * layers
+        assert len(block_types) == layers
+
+        def _create_block(bt: str):
+            if bt == 'CustomResidualAttentionBlock':
+                return CustomResidualAttentionBlock(
+                    width,
+                    heads,
+                    mlp_ratio=mlp_ratio,
+                    ls_init_value=ls_init_value,
+                    act_layer=act_layer,
+                    norm_layer=norm_layer,
+                    batch_first=batch_first,
+                )
+            else:
+                assert False
+
+        self.resblocks = nn.ModuleList([
+            _create_block(bt)
+            for bt in block_types
+        ])
+
+    def get_cast_dtype(self) -> torch.dtype:
+        return self.resblocks[0].get_weight_dtype()
+
+    def forward_intermediates(
+            self,
+            x: torch.Tensor,
+            attn_mask: Optional[torch.Tensor] = None,
+            indices: Optional[Union[int, List[int]]] = None,
+            stop_early: bool = False,
+    ):
+        take_indices, max_index = feature_take_indices(len(self.resblocks), indices)
+
+        if not self.batch_first:
+            x = x.transpose(0, 1).contiguous()  # NLD -> LND
+
+        intermediates = []
+        if torch.jit.is_scripting() or not stop_early:  # can't slice blocks in torchscript
+            blocks = self.resblocks
+        else:
+            blocks = self.resblocks[:max_index + 1]
+        for i, blk in enumerate(blocks):
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                x = checkpoint(blk, x, None, None, attn_mask, use_reentrant=False)
+            else:
+                x = blk(x, attn_mask=attn_mask)
+
+            if i in take_indices:
+                intermediates.append(x.transpose(0, 1) if not self.batch_first else x)
+
+        if not self.batch_first:
+            x = x.transpose(0, 1)  # LND -> NLD
+
+        return x, intermediates
+
+    def prune_intermediate_layers(self, indices: Union[int, List[int]] = 1):
+        """ Prune layers not required for specified intermediates.
+        """
+        take_indices, max_index = feature_take_indices(len(self.resblocks), indices)
+        self.resblocks = self.resblocks[:max_index + 1]  # truncate blocks
+        return take_indices
+
+    def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
+        if not self.batch_first:
+            x = x.transpose(0, 1)  # NLD -> LND
+
+        for r in self.resblocks:
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                # TODO: handle kwargs https://github.com/pytorch/pytorch/issues/79887#issuecomment-1161758372
+                x = checkpoint(r, x, None, None, attn_mask, use_reentrant=False)
+            else:
+                x = r(x, attn_mask=attn_mask)
+
+        if not self.batch_first:
+            x = x.transpose(0, 1)  # NLD -> LND
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(
+            self,
+            width: int,
+            layers: int,
+            heads: int,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Type[nn.Module] = nn.GELU,
+            norm_layer: Type[nn.Module] = LayerNorm,
+            batch_first: bool = True,
+            block_type: Optional[str] = None,
+            qk_norm: bool = False,
+            scaled_cosine_attn: bool = False,
+            scale_heads: bool = False,
+            scale_attn_inner: bool = False,
+            scale_attn: bool = False,
+            scale_fc: bool = False,
+    ):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.batch_first = batch_first
+        self.grad_checkpointing = False
+
+        # Auto-select custom block if any custom features are enabled
+        if block_type is None:
+            if any([qk_norm, scaled_cosine_attn, scale_heads, scale_attn_inner, scale_attn, scale_fc]):
+                block_type = 'custom'
+            else:
+                block_type = 'default'
+
+        if block_type == 'custom':
+            self.resblocks = nn.ModuleList([
+                CustomResidualAttentionBlock(
+                    width,
+                    heads,
+                    mlp_ratio,
+                    ls_init_value=ls_init_value,
+                    act_layer=act_layer,
+                    norm_layer=norm_layer,
+                    qk_norm=qk_norm,
+                    scale_cosine_attn=scaled_cosine_attn,
+                    scale_heads=scale_heads,
+                    scale_attn_inner=scale_attn_inner,
+                    scale_attn=scale_attn,
+                    scale_fc=scale_fc,
+                    batch_first=batch_first,
+                )
+                for _ in range(layers)
+            ])
+        else:
+            self.resblocks = nn.ModuleList([
+                ResidualAttentionBlock(
+                    width,
+                    heads,
+                    mlp_ratio,
+                    ls_init_value=ls_init_value,
+                    act_layer=act_layer,
+                    norm_layer=norm_layer,
+                    batch_first=batch_first,
+                )
+                for _ in range(layers)
+            ])
+
+    def get_cast_dtype(self) -> torch.dtype:
+        return self.resblocks[0].get_weight_dtype()
+
+    def forward_intermediates(
+            self,
+            x: torch.Tensor,
+            attn_mask: Optional[torch.Tensor] = None,
+            indices: Optional[Union[int, List[int]]] = None,
+            stop_early: bool = False,
+    ):
+        take_indices, max_index = feature_take_indices(len(self.resblocks), indices)
+
+        if not self.batch_first:
+            x = x.transpose(0, 1).contiguous()    # NLD -> LND
+
+        intermediates = []
+        if torch.jit.is_scripting() or not stop_early:  # can't slice blocks in torchscript
+            blocks = self.resblocks
+        else:
+            blocks = self.resblocks[:max_index + 1]
+        for i, blk in enumerate(blocks):
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                x = checkpoint(blk, x, None, None, attn_mask, use_reentrant=False)
+            else:
+                x = blk(x, attn_mask=attn_mask)
+
+            if i in take_indices:
+                intermediates.append(x.transpose(0, 1) if not self.batch_first else x)
+
+        if not self.batch_first:
+            x = x.transpose(0, 1)    # LND -> NLD
+
+        return x, intermediates
+
+    def prune_intermediate_layers(self, indices: Union[int, List[int]] = 1):
+        """ Prune layers not required for specified intermediates.
+        """
+        take_indices, max_index = feature_take_indices(len(self.resblocks), indices)
+        self.resblocks = self.resblocks[:max_index + 1]  # truncate blocks
+        return take_indices
+
+    def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
+        if not self.batch_first:
+            x = x.transpose(0, 1).contiguous()    # NLD -> LND
+
+        for r in self.resblocks:
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                # TODO: handle kwargs https://github.com/pytorch/pytorch/issues/79887#issuecomment-1161758372
+                x = checkpoint(r, x, None, None, attn_mask, use_reentrant=False)
+            else:
+                x = r(x, attn_mask=attn_mask)
+
+        if not self.batch_first:
+            x = x.transpose(0, 1)    # LND -> NLD
+        return x
+
+
+def _expand_token(token, batch_size: int):
+    return token.view(1, 1, -1).expand(batch_size, -1, -1)
+
+
+class VisionTransformer(nn.Module):
+    output_tokens: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            image_size: int,
+            patch_size: int,
+            width: int,
+            layers: int,
+            heads: int,
+            mlp_ratio: float,
+            ls_init_value: float = None,
+            attentional_pool: bool = False,
+            attn_pooler_queries: int = 256,
+            attn_pooler_heads: int = 8,
+            output_dim: int = 512,
+            patch_dropout: float = 0.,
+            no_ln_pre: bool = False,
+            pos_embed_type: str = 'learnable',
+            pool_type: str = 'tok',
+            final_ln_after_pool: bool = False,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            output_tokens: bool = False,
+            block_type: Optional[str] = None,
+            qk_norm: bool = False,
+            scaled_cosine_attn: bool = False,
+            scale_heads: bool = False,
+            scale_attn_inner: bool = False,
+            scale_attn: bool = False,
+            scale_fc: bool = False,
+    ):
+        super().__init__()
+        assert pool_type in ('tok', 'avg', 'none')
+        self.output_tokens = output_tokens
+        image_height, image_width = self.image_size = to_2tuple(image_size)
+        patch_height, patch_width = self.patch_size = to_2tuple(patch_size)
+        self.grid_size = (image_height // patch_height, image_width // patch_width)
+        self.final_ln_after_pool = final_ln_after_pool  # currently ignored w/ attn pool enabled
+        self.output_dim = output_dim
+
+        self.conv1 = nn.Conv2d(
+            in_channels=3,
+            out_channels=width,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=False,
+        )
+
+        # class embeddings and positional embeddings
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        if pos_embed_type == 'learnable':
+            self.positional_embedding = nn.Parameter(
+                scale * torch.randn(self.grid_size[0] * self.grid_size[1] + 1, width))
+        elif pos_embed_type == 'sin_cos_2d':
+            # fixed sin-cos embedding
+            assert self.grid_size[0] == self.grid_size[1],\
+                'currently sin cos 2d pos embedding only supports square input'
+            self.positional_embedding = nn.Parameter(
+                torch.zeros(self.grid_size[0] * self.grid_size[1] + 1, width), requires_grad=False)
+            pos_embed_type = get_2d_sincos_pos_embed(width, self.grid_size[0], cls_token=True)
+            self.positional_embedding.data.copy_(torch.from_numpy(pos_embed_type).float())
+        else:
+            raise ValueError
+
+        # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
+        self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()
+
+        self.ln_pre = nn.Identity() if no_ln_pre else norm_layer(width)
+        self.transformer = Transformer(
+            width,
+            layers,
+            heads,
+            mlp_ratio,
+            ls_init_value=ls_init_value,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            block_type=block_type,
+            qk_norm=qk_norm,
+            scaled_cosine_attn=scaled_cosine_attn,
+            scale_heads=scale_heads,
+            scale_attn_inner=scale_attn_inner,
+            scale_attn=scale_attn,
+            scale_fc=scale_fc,
+        )
+
+        if attentional_pool:
+            if isinstance(attentional_pool, str):
+                self.attn_pool_type = attentional_pool
+                self.pool_type = 'none'
+                if attentional_pool in ('parallel', 'cascade'):
+                    self.attn_pool = AttentionalPooler(
+                        output_dim,
+                        width,
+                        n_head=attn_pooler_heads,
+                        n_queries=attn_pooler_queries,
+                    )
+                    self.attn_pool_contrastive = AttentionalPooler(
+                        output_dim,
+                        width,
+                        n_head=attn_pooler_heads,
+                        n_queries=1,
+                    )
+                else:
+                    assert False
+            else:
+                self.attn_pool_type = ''
+                self.pool_type = pool_type
+                self.attn_pool = AttentionalPooler(
+                    output_dim,
+                    width,
+                    n_head=attn_pooler_heads,
+                    n_queries=attn_pooler_queries,
+                )
+                self.attn_pool_contrastive = None
+            pool_dim = output_dim
+        else:
+            self.attn_pool = None
+            pool_dim = width
+            self.pool_type = pool_type
+
+        self.ln_post = norm_layer(pool_dim)
+        self.proj = nn.Parameter(scale * torch.randn(pool_dim, output_dim))
+
+        self.init_parameters()
+
+    def lock(self, unlocked_groups: int = 0, freeze_bn_stats: bool = False):
+        for param in self.parameters():
+            param.requires_grad = False
+
+        if unlocked_groups != 0:
+            groups = [
+                [
+                    self.conv1,
+                    self.class_embedding,
+                    self.positional_embedding,
+                    self.ln_pre,
+                ],
+                *self.transformer.resblocks[:-1],
+                [
+                    self.transformer.resblocks[-1],
+                    self.ln_post,
+                ],
+                self.proj,
+            ]
+
+            def _unlock(x):
+                if isinstance(x, Sequence):
+                    for g in x:
+                        _unlock(g)
+                else:
+                    if isinstance(x, torch.nn.Parameter):
+                        x.requires_grad = True
+                    else:
+                        for p in x.parameters():
+                            p.requires_grad = True
+
+            _unlock(groups[-unlocked_groups:])
+
+    def init_parameters(self):
+        # FIXME OpenAI CLIP did not define an init for the VisualTransformer
+        # TODO experiment if default PyTorch init, below, or alternate init is best.
+
+        # nn.init.normal_(self.class_embedding, std=self.scale)
+        # nn.init.normal_(self.positional_embedding, std=self.scale)
+        #
+        # 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.scale)
+        pass
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable: bool = True):
+        self.transformer.grad_checkpointing = enable
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        # for timm optimizers, 1d params like logit_scale, logit_bias, ln/bn scale, biases are excluded by default
+        no_wd = {'positional_embedding', 'class_embedding'}
+        return no_wd
+
+    def _global_pool(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.pool_type == 'avg':
+            pooled, tokens = x[:, 1:].mean(dim=1), x[:, 1:]
+        elif self.pool_type == 'tok':
+            pooled, tokens = x[:, 0], x[:, 1:]
+        else:
+            pooled = tokens = x
+
+        return pooled, tokens
+
+    def _embeds(self, x:torch.Tensor) -> torch.Tensor:
+        x = self.conv1(x)  # shape = [*, dim, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+
+        # class embeddings and positional embeddings
+        x = torch.cat([_expand_token(self.class_embedding, x.shape[0]).to(x.dtype), x], dim=1)
+        # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+
+        # patch dropout (if active)
+        x = self.patch_dropout(x)
+
+        # apply norm before transformer
+        x = self.ln_pre(x)
+        return x
+
+    def _pool(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.attn_pool is not None:
+            if self.attn_pool_contrastive is not None:
+                # This is untested, WIP pooling that should match paper
+                x = self.ln_post(x)  # TBD LN first or separate one after each pool?
+                tokens = self.attn_pool(x)
+                if self.attn_pool_type == 'parallel':
+                    pooled = self.attn_pool_contrastive(x)
+                else:
+                    assert self.attn_pool_type == 'cascade'
+                    pooled = self.attn_pool_contrastive(tokens)
+            else:
+                # this is the original OpenCLIP CoCa setup, does not match paper
+                x = self.attn_pool(x)
+                x = self.ln_post(x)
+                pooled, tokens = self._global_pool(x)
+        elif self.final_ln_after_pool:
+            pooled, tokens = self._global_pool(x)
+            pooled = self.ln_post(pooled)
+        else:
+            x = self.ln_post(x)
+            pooled, tokens = self._global_pool(x)
+
+        return pooled, tokens
+
+    def forward_intermediates(
+            self,
+            x: torch.Tensor,
+            indices: Optional[Union[int, List[int]]] = None,
+            stop_early: bool = False,
+            normalize_intermediates: bool = False,
+            intermediates_only: bool = False,
+            output_fmt: str = 'NCHW',
+            output_extra_tokens: bool = False,
+    ) -> Dict[str, Union[torch.Tensor, List[torch.Tensor]]]:
+        """ Forward features that returns intermediates.
+
+        Args:
+            x: Input image tensor
+            indices: Take last n blocks if int, all if None, select matching indices if sequence
+            stop_early: Stop iterating over blocks when last desired intermediate hit
+            intermediates_only: Only return intermediate features
+            normalize_intermediates: Apply final norm layer to all intermediates
+            output_fmt: Shape of intermediate feature outputs
+            output_extra_tokens: Return both extra prefix class tokens
+        Returns:
+
+        """
+        assert output_fmt in ('NCHW', 'NLC'), 'Output format must be one of NCHW or NLC.'
+        reshape = output_fmt == 'NCHW'
+
+        # forward pass
+        B, _, height, width = x.shape
+        x = self._embeds(x)
+        x, intermediates = self.transformer.forward_intermediates(
+            x,
+            indices=indices,
+            stop_early=stop_early,
+        )
+
+        # process intermediates
+        if normalize_intermediates:
+            # apply final norm to all intermediates
+            intermediates = [self.ln_post(xi) for xi in intermediates]
+        num_prefix_tokens = 1  # one class token that's always there (as of now)
+        if num_prefix_tokens:
+            # split prefix (e.g. class, distill) and spatial feature tokens
+            prefix_tokens = [y[:, 0:num_prefix_tokens] for y in intermediates]
+            intermediates = [y[:, num_prefix_tokens:] for y in intermediates]
+        else:
+            prefix_tokens = None
+        if reshape:
+            # reshape to BCHW output format
+            H, W = height // self.patch_size[0], width // self.patch_size[1]
+            intermediates = [y.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() for y in intermediates]
+
+        output = {'image_intermediates': intermediates}
+        if prefix_tokens is not None and output_extra_tokens:
+            output['image_intermediates_prefix'] = prefix_tokens
+
+        if intermediates_only:
+            return output
+
+        pooled, _ = self._pool(x)
+
+        if self.proj is not None:
+            pooled = pooled @ self.proj
+
+        output['image_features'] = pooled
+
+        return output
+
+    def prune_intermediate_layers(
+            self,
+            indices: Union[int, List[int]] = 1,
+            prune_norm: bool = False,
+            prune_head: bool = True,
+    ):
+        """ Prune layers not required for specified intermediates.
+        """
+        take_indices = self.transformer.prune_intermediate_layers(indices)
+        if prune_norm:
+            self.ln_post = nn.Identity()
+        if prune_head:
+            self.proj = None
+        return take_indices
+
+    def forward(self, x: torch.Tensor):
+        x = self._embeds(x)
+        x = self.transformer(x)
+        pooled, tokens = self._pool(x)
+
+        if self.proj is not None:
+            pooled = pooled @ self.proj
+
+        if self.output_tokens:
+            return pooled, tokens
+        
+        return pooled
+
+
+def text_global_pool(
+        x: torch.Tensor,
+        text: Optional[torch.Tensor] = None,
+        pool_type: str = 'argmax',
+        eos_token_id: Optional[int] = None,
+) -> torch.Tensor:
+    if pool_type == 'first':
+        pooled = x[:, 0]
+    elif pool_type == 'last':
+        pooled = x[:, -1]
+    elif pool_type == 'argmax':
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        assert text is not None
+        pooled = x[torch.arange(x.shape[0], device=x.device), text.argmax(dim=-1)]
+    elif pool_type == 'eos':
+        # take features from tokenizer specific eos
+        assert text is not None
+        assert eos_token_id is not None
+        idx = (text == eos_token_id).int().argmax(dim=-1)
+        pooled = x[torch.arange(x.shape[0], device=x.device), idx]
+    else:
+        pooled = x
+
+    return pooled
+
+
+class TextTransformer(nn.Module):
+    output_tokens: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            context_length: int = 77,
+            vocab_size: int = 49408,
+            width: int = 512,
+            heads: int = 8,
+            layers: int = 12,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            output_dim: Optional[int] = 512,
+            embed_cls: bool = False,
+            no_causal_mask: bool = False,
+            use_pad_mask: bool = False,
+            correct_cls_mask: bool = False,
+            pad_id: int = 0,
+            eos_id: int = 2,
+            pool_type: str = 'argmax',
+            proj_type: str = 'linear',
+            proj_bias: bool = False,
+            act_layer: Type[nn.Module] = nn.GELU,
+            norm_layer: Type[nn.Module] = LayerNorm,
+            output_tokens: bool = False,
+            block_type: Optional[str] = None,
+            qk_norm: bool = False,
+            scaled_cosine_attn: bool = False,
+            scale_heads: bool = False,
+            scale_attn_inner: bool = False,
+            scale_attn: bool = False,
+            scale_fc: bool = False,
+    ):
+        super().__init__()
+        assert pool_type in ('first', 'last', 'argmax', 'eos', 'none')
+        self.output_tokens = output_tokens
+        self.num_pos = self.context_length = context_length
+        self.vocab_size = vocab_size
+        self.width = width
+        self.output_dim = output_dim
+        self.heads = heads
+        self.pad_id = pad_id
+        self.eos_id = eos_id
+        self.pool_type = pool_type
+        self.use_pad_mask = use_pad_mask and no_causal_mask  # only use in bi‑dir mode
+        self.correct_cls_mask = correct_cls_mask  # use the correct cls mask for CoCa (original is wrong)
+
+        self.token_embedding = nn.Embedding(vocab_size, width)
+        if embed_cls:
+            self.cls_emb = nn.Parameter(torch.empty(width))
+            self.num_pos += 1
+        else:
+            self.cls_emb = None
+        self.positional_embedding = nn.Parameter(torch.empty(self.num_pos, width))
+        self.transformer = Transformer(
+            width=width,
+            layers=layers,
+            heads=heads,
+            mlp_ratio=mlp_ratio,
+            ls_init_value=ls_init_value,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            block_type=block_type,
+            qk_norm=qk_norm,
+            scaled_cosine_attn=scaled_cosine_attn,
+            scale_heads=scale_heads,
+            scale_attn_inner=scale_attn_inner,
+            scale_attn=scale_attn,
+            scale_fc=scale_fc,
+        )
+        self.ln_final = norm_layer(width)
+
+        if no_causal_mask:
+            self.attn_mask = None  # bi‑directional
+        else:
+            self.register_buffer('attn_mask', self.build_causal_mask(), persistent=False)
+
+        if proj_type == 'none' or not output_dim:
+            self.text_projection = None
+        else:
+            if proj_bias:
+                self.text_projection = nn.Linear(width, output_dim)
+            else:
+                self.text_projection = nn.Parameter(torch.empty(width, output_dim))
+
+        self.init_parameters()
+
+    def init_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+        if self.cls_emb is not None:
+            nn.init.normal_(self.cls_emb, std=0.01)
+
+        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:
+            if isinstance(self.text_projection, nn.Linear):
+                nn.init.normal_(self.text_projection.weight, std=self.transformer.width ** -0.5)
+                if self.text_projection.bias is not None:
+                    nn.init.zeros_(self.text_projection.bias)
+            else:
+                nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.transformer.grad_checkpointing = enable
+
+    def lock(self, unlocked_layers: int = 0, freeze_layer_norm: bool = True):
+        """
+        Lock the text transformer layers, optionally leaving some layers unlocked.
+
+        Args:
+            unlocked_layers: Number of layers to leave unlocked (from the end).
+            freeze_layer_norm: LayerNorm freeze (only for API compatibility, not functional)
+        """
+        assert freeze_layer_norm, 'Unfreezing LayerNorm is not supported. LayerNorm treated like other weights.'
+        lock_text_tower(self, unlocked_layers)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        # for timm optimizers, 1d params like logit_scale, logit_bias, ln/bn scale, biases are excluded by default
+        no_wd = {'positional_embedding'}
+        if self.cls_emb is not None:
+            no_wd.add('cls_emb')
+        return no_wd
+
+    def build_causal_mask(self):
+        # lazily create causal attention mask, with full attention between the tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.num_pos, self.num_pos)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    def _build_additive_mask(
+        self,
+        text: torch.Tensor,  # [B, L] – original text ids without CLS yet
+        seq_len: int,  # L (+1 if CLS added)
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """
+        Returns an additive (-inf) mask of shape [B*heads, seq_len, seq_len] that
+        simultaneously masks padding tokens and (optionally) the CLS token.
+        """
+        valid = text != self.pad_id  # [B, L] (True = keep)
+
+        if self.cls_emb is not None:
+            cls_valid = valid.new_ones(valid.size(0), 1) # [B, 1]
+            # cls mask pos at end if correct or front for incorrect legacy mode in existing CoCa weights
+            valid = torch.cat([valid, cls_valid] if self.correct_cls_mask else [cls_valid, valid], 1)
+
+        # broadcast over query dimension
+        key_mask = valid.unsqueeze(1).expand(-1, seq_len, -1)  # [B, Q, K]
+        additive = torch.zeros_like(key_mask, dtype=dtype)
+        additive.masked_fill_(~key_mask, float("-inf"))
+        additive = additive.repeat_interleave(self.heads, 0)  # [B*H, Q, K]
+        return additive
+
+    def _embeds(self, text) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        cast_dtype = self.transformer.get_cast_dtype()
+        B, seq_len = text.shape
+
+        x = self.token_embedding(text).to(cast_dtype)
+
+        # Optional class token (always appended ala CoCa)
+        if self.cls_emb is not None:
+            x = torch.cat([x, _expand_token(self.cls_emb, x.size(0))], 1)
+            seq_len += 1
+
+        attn_mask = self.attn_mask  # Base causal mask (if any)
+
+        # Class + padding additive mask
+        if self.use_pad_mask or self.cls_emb is not None:
+            add_mask  = self._build_additive_mask(text, seq_len, x.dtype)
+            if attn_mask is not None:
+                # Slice the causal mask to match current sequence length
+                attn_mask = attn_mask[:seq_len, :seq_len].unsqueeze(0) + add_mask
+            else:
+                attn_mask = add_mask
+
+        x = x + self.positional_embedding[:seq_len].to(cast_dtype)
+        return x, attn_mask
+
+    def forward_intermediates(
+            self,
+            text: torch.Tensor,
+            indices: Optional[Union[int, List[int]]] = None,
+            stop_early: bool = False,
+            normalize_intermediates: bool = False,
+            intermediates_only: bool = False,
+            output_fmt: str = 'NCHW',
+            output_extra_tokens: bool = False,
+    ) -> Dict[str, Union[torch.Tensor, List[torch.Tensor]]]:
+        """ Forward features that returns intermediates.
+
+        Args:
+            text: Input text ids
+            indices: Take last n blocks if int, all if None, select matching indices if sequence
+            stop_early: Stop iterating over blocks when last desired intermediate hit
+            normalize_intermediates: Apply norm layer to all intermediates
+            intermediates_only: Only return intermediate features
+            output_fmt: Shape of intermediate feature outputs
+            output_extra_tokens: Return both prefix and intermediate tokens
+        Returns:
+
+        """
+        assert output_fmt in ('NLC',), 'Output format must be NLC.'
+        # forward pass
+        x, attn_mask = self._embeds(text)
+        x, intermediates = self.transformer.forward_intermediates(
+            x,
+            attn_mask=attn_mask,
+            indices=indices,
+            stop_early=stop_early,
+        )
+
+        # process intermediates
+        if normalize_intermediates:
+            # apply final norm to all intermediates
+            intermediates = [self.ln_final(xi) for xi in intermediates]
+
+        output = {}
+
+        if self.cls_emb is not None:
+            seq_intermediates = [xi[:, :-1] for xi in intermediates]  # separate concat'd class token from sequence
+            if output_extra_tokens:
+                # return suffix class tokens separately
+                cls_intermediates = [xi[:, -1:] for xi in intermediates]
+                output['text_intermediates_suffix'] = cls_intermediates
+            intermediates = seq_intermediates
+        output['text_intermediates'] = intermediates
+
+        if intermediates_only:
+            return output
+
+        if self.cls_emb is not None:
+            # presence of appended cls embed (CoCa) overrides pool_type, always take last token
+            pooled = text_global_pool(x, pool_type='last')
+            pooled = self.ln_final(pooled)  # final LN applied after pooling in this case
+        else:
+            x = self.ln_final(x)
+            pooled = text_global_pool(x, text, pool_type=self.pool_type, eos_token_id=getattr(self, "eos_id", None))
+
+        if self.text_projection is not None:
+            if isinstance(self.text_projection, nn.Linear):
+                pooled = self.text_projection(pooled)
+            else:
+                pooled = pooled @ self.text_projection
+
+        output['text_features'] = pooled
+
+        return output
+
+    def prune_intermediate_layers(
+            self,
+            indices: Union[int, List[int]] = 1,
+            prune_norm: bool = False,
+            prune_head: bool = True,
+    ):
+        """ Prune layers not required for specified intermediates.
+        """
+        take_indices = self.transformer.prune_intermediate_layers(indices)
+        if prune_norm:
+            self.ln_final = nn.Identity()
+        if prune_head:
+            self.text_projection = None
+        return take_indices
+
+    def forward(self, text):
+        x, attn_mask = self._embeds(text)
+
+        x = self.transformer(x, attn_mask=attn_mask)
+
+        # x.shape = [batch_size, n_ctx, transformer.width]
+        if self.cls_emb is not None:
+            # presence of appended cls embed (CoCa) overrides pool_type, always take last token
+            pooled = text_global_pool(x, pool_type='last')
+            pooled = self.ln_final(pooled)  # final LN applied after pooling in this case
+            tokens = x[:, :-1]
+        else:
+            x = self.ln_final(x)
+            pooled = text_global_pool(x, text, pool_type=self.pool_type, eos_token_id=getattr(self, "eos_id", None))
+            tokens = x
+
+        if self.text_projection is not None:
+            if isinstance(self.text_projection, nn.Linear):
+                pooled = self.text_projection(pooled)
+            else:
+                pooled = pooled @ self.text_projection
+
+        if self.output_tokens:
+            return pooled, tokens
+
+        return pooled
+
+
+class MultimodalTransformer(Transformer):
+    """Cross-attention based multimodal decoder.
+    
+    Text and image/biosignals embeddings are kept separate.
+    Each layer has:
+      1. Self-attention on text (causal)
+      2. Cross-attention from text to image/biosignals
+    """
+    def __init__(
+            self,
+            width: int,
+            layers: int,
+            heads: int,
+            context_length: int = 77,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Type[nn.Module] = nn.GELU,
+            norm_layer: Type[nn.Module] = LayerNorm,
+            output_dim: int = 512,
+            batch_first: bool = True,
+            prefix_len: int = 0,
+    ):
+        super().__init__(
+            width=width,
+            layers=layers,
+            heads=heads,
+            mlp_ratio=mlp_ratio,
+            ls_init_value=ls_init_value,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            batch_first=batch_first,
+        )
+        self.context_length = context_length
+        self.cross_attn = nn.ModuleList([
+            ResidualAttentionBlock(
+                width,
+                heads,
+                mlp_ratio,
+                ls_init_value=ls_init_value,
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                is_cross_attention=True,
+                batch_first=batch_first,
+            )
+            for _ in range(layers)
+        ])
+
+        # Register attention masks based on prefix configuration
+        self.prefix_len = prefix_len
+        if prefix_len > 0:
+            # Pre-build prefix-causal mask for condition tokens + text
+            prefix_causal_mask = self.build_prefix_causal_mask(prefix_len, context_length)
+            self.register_buffer('prefix_causal_mask', prefix_causal_mask, persistent=False)
+        else:
+            # Only register standard causal mask when not using prefix tokens
+            self.register_buffer('attn_mask', self.build_attention_mask(), persistent=False)
+
+        self.ln_final = norm_layer(width)
+        self.text_projection = nn.Parameter(torch.empty(width, output_dim))
+        
+        self.init_parameters()
+
+    def init_parameters(self):
+        proj_std = (self.width ** -0.5) * ((2 * self.layers) ** -0.5)
+        attn_std = self.width ** -0.5
+        fc_std = (2 * self.width) ** -0.5
+        for block in self.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)
+        for block in self.cross_attn:
+            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.width ** -0.5)
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the 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
+
+    # def build_prefix_causal_mask(self, prefix_len: int, text_len: int):
+    #     """Build a prefix-causal attention mask for condition tokens + text.
+        
+    #     Args:
+    #         prefix_len: Length of prefix (condition tokens)
+    #             These tokens receive full bidirectional attention among themselves.
+    #         text_len: Length of text sequence
+    #             These tokens receive causal attention.
+            
+    #     Returns:
+    #         Additive mask of shape (prefix_len + text_len, prefix_len + text_len)
+    #         Where -inf = cannot attend, 0 = can attend
+            
+    #     Attention pattern:
+    #         - Prefix tokens ↔ Prefix tokens: Full bidirectional (can attend)
+    #         - Text tokens → Prefix tokens: Full attention (can attend)
+    #         - Text tokens → Text tokens: Causal attention (only previous tokens)
+    #         - Prefix tokens → Text tokens: Cannot attend (masked)
+    #     """
+    #     total_len = prefix_len + text_len
+    #     mask = torch.zeros(total_len, total_len)
+        
+    #     # Prefix tokens can attend to all prefix tokens (bidirectional)
+    #     # mask[:prefix_len, :prefix_len] remains 0 (can attend)
+        
+    #     # Prefix tokens cannot attend to text tokens
+    #     mask[:prefix_len, prefix_len:] = float("-inf")
+        
+    #     # Text tokens can attend to all prefix tokens
+    #     # mask[prefix_len:, :prefix_len] remains 0 (can attend)
+        
+    #     # Text tokens attend to previous text tokens only (causal)
+    #     text_causal_mask = torch.triu(torch.ones(text_len, text_len), diagonal=1) * float("-inf")
+    #     mask[prefix_len:, prefix_len:] = text_causal_mask
+        
+    #     return mask
+
+    def build_prefix_causal_mask(self, prefix_len: int, text_len: int):
+        """Additive mask; 0 = attend, NEG = block (fp32 for stability)."""
+        total_len = prefix_len + text_len
+        # fp32 on CPU; we'll .to(device) later without changing dtype
+        mask = torch.zeros(total_len, total_len, dtype=torch.float32)
+
+        # large finite negative (safer than -inf for fp16/bf16 kernels)
+        NEG = -torch.finfo(mask.dtype).max
+
+        # Prefix → Text: block
+        mask[:prefix_len, prefix_len:] = NEG
+
+        # Text → Text: causal (block future). Use masked_fill, not 0 * -inf.
+        tri = torch.triu(torch.ones(text_len, text_len, dtype=torch.bool), diagonal=1)
+        mask[prefix_len:, prefix_len:].masked_fill_(tri, NEG)
+        return mask
+
+    def forward_intermediates(
+            self,
+            x: torch.Tensor,
+            attn_mask: Optional[torch.Tensor] = None,
+            indices: Optional[Union[int, List[int]]] = None,
+            stop_early: bool = False,
+    ):
+        assert False, "Not currently implemented for MultimodalTransformer w/ xattn"
+
+    def forward(self, image_embs, text_embs, condition_embs=None):
+        """Forward pass with cross-attention between text and image.
+        
+        Args:
+            image_embs: (batch_size, num_image_tokens, width)
+            text_embs: (batch_size, num_text_tokens, width)
+            condition_embs: Optional (batch_size, num_condition_tokens, width)
+                Additional conditioning tokens that will be prepended to text.
+                These tokens get full bidirectional attention among themselves,
+                then cross-attend to image embeddings.
+            
+        Returns:
+            Text decoder outputs: (batch_size, num_text_tokens, output_dim)
+                Note: Only text token outputs are returned (condition token outputs are excluded)
+        """
+        # Determine text length before prepending conditions
+        original_text_len = text_embs.shape[1]
+        assert original_text_len <= self.context_length, "original_text_len must be less than or equal to context_length"
+        
+        # Prepend condition tokens to text if provided
+        if condition_embs is not None:
+            condition_len = condition_embs.shape[1]
+            
+            # Safety check: condition_len must not exceed the pre-configured prefix_len
+            assert condition_len <= self.prefix_len, \
+                f"condition_len ({condition_len}) exceeds prefix_len ({self.prefix_len})"
+            
+            text_embs = torch.cat([condition_embs, text_embs], dim=1)  # (batch, cond_len + text_len, width)
+        else:
+            condition_len = 0
+        
+        # Get attention mask based on prefix configuration
+        if self.prefix_len > 0:
+            # Slice the pre-built prefix-causal mask based on actual condition_len
+            # The mask is built for (prefix_len + context_length)
+            # When condition_len < prefix_len, we slice from offset to get the right structure
+            offset = self.prefix_len - condition_len  # How many prefix positions to skip
+            seq_len = condition_len + original_text_len  # Total sequence length
+            attn_mask = self.prefix_causal_mask[offset:offset+seq_len, offset:offset+seq_len].to(device=text_embs.device)
+        else:
+            # Use standard causal mask when prefix_len == 0
+            seq_len = original_text_len
+            attn_mask = self.attn_mask[:seq_len, :seq_len].to(device=text_embs.device)
+        
+        if not self.batch_first:
+            image_embs = image_embs.permute(1, 0, 2)  # NLD -> LND
+            text_embs = text_embs.permute(1, 0, 2)  # NLD -> LND
+
+        for resblock, cross_attn in zip(self.resblocks, self.cross_attn):
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                # TODO: handle kwargs https://github.com/pytorch/pytorch/issues/79887#issuecomment-1161758372
+                text_embs = checkpoint(
+                    resblock, text_embs, None, None, attn_mask, use_reentrant=False)
+                text_embs = checkpoint(
+                    cross_attn, text_embs, image_embs, image_embs, None, use_reentrant=False)
+            else:
+                text_embs = resblock(text_embs, attn_mask=attn_mask)
+                text_embs = cross_attn(text_embs, k_x=image_embs, v_x=image_embs)
+
+        if not self.batch_first:
+            text_embs = text_embs.permute(1, 0, 2)  # LND -> NLD
+
+        out = self.ln_final(text_embs)
+        if self.text_projection is not None:
+            out = out @ self.text_projection
+        
+        # Extract only the text portion (skip condition tokens if present)
+        if condition_len > 0:
+            out = out[:, condition_len:, :]  # (batch, text_len, output_dim)
+
+        return out
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.grad_checkpointing = enable
+
+
+class ConcatMultimodalTransformer(Transformer):
+    """Concatenation-based multimodal decoder.
+    
+    Concatenates [condition_tokens (optional), image/biosignals_tokens, text_tokens] into a single sequence.
+    Uses unified self-attention with a prefix-causal mask that allows:
+      - Condition tokens attend to all condition + image tokens (full bidirectional)
+      - Image/biosignals tokens attend to all condition + image tokens (full bidirectional)
+      - Text tokens attend to all condition + image tokens (full attention to prefix)
+      - Text tokens attend to all previous text tokens (causal)
+    
+    This enables flexible conditioning where any prefix tokens (condition + image) get full
+    bidirectional attention, while text maintains causal generation properties.
+    """
+    def __init__(
+            self,
+            width: int,
+            layers: int,
+            heads: int,
+            context_length: int = 77,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Type[nn.Module] = nn.GELU,
+            norm_layer: Type[nn.Module] = LayerNorm,
+            output_dim: int = 512,
+            batch_first: bool = True,
+            prefix_len: int = 0,
+    ):
+        super().__init__(
+            width=width,
+            layers=layers,
+            heads=heads,
+            mlp_ratio=mlp_ratio,
+            ls_init_value=ls_init_value,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            batch_first=batch_first,
+        )
+        self.context_length = context_length
+        self.condition_prefix_len = prefix_len  # Number of condition tokens (0, 1, or N)
+        
+        # Pre-register an empty buffer for the attention mask
+        # Will be populated on first forward pass when image token count is known
+        self.register_buffer('_cached_attn_mask', torch.empty(0), persistent=False)
+        self._cached_prefix_len = None  # Track the prefix length used to build the cache
+        
+        # No cross-attention layers needed - uses self-attention only
+        self.ln_final = norm_layer(width)
+        self.text_projection = nn.Parameter(torch.empty(width, output_dim))
+
+        # self.init_parameters()
+
+    def init_parameters(self):
+        proj_std = (self.width ** -0.5) * ((2 * self.layers) ** -0.5)
+        attn_std = self.width ** -0.5
+        fc_std = (2 * self.width) ** -0.5
+        for block in self.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.width ** -0.5)
+
+    # def build_prefix_causal_mask(self, prefix_len: int, text_len: int):
+    #     """Build a prefix-causal attention mask.
+        
+    #     Args:
+    #         prefix_len: Length of the prefix (condition + image/biosignals tokens)
+    #             All prefix tokens receive full bidirectional attention among themselves.
+    #         text_len: Length of text sequence
+            
+    #     Returns:
+    #         Additive mask of shape (prefix_len + text_len, prefix_len + text_len)
+    #         Where -inf = cannot attend, 0 = can attend
+            
+    #     Attention pattern:
+    #         - Prefix tokens ↔ Prefix tokens: Full bidirectional (can attend)
+    #         - Text tokens → Prefix tokens: Full attention (can attend)
+    #         - Text tokens → Text tokens: Causal attention (only previous tokens)
+    #         - Prefix tokens → Text tokens: Cannot attend (masked)
+    #     """
+    #     total_len = prefix_len + text_len
+    #     # Start with a float mask of zeros (all positions can attend)
+    #     mask = torch.zeros(total_len, total_len, dtype=torch.float32)
+        
+    #     # Prefix tokens can attend to all prefix tokens (bidirectional)
+    #     # mask[:prefix_len, :prefix_len] remains 0 (can attend)
+        
+    #     # Prefix tokens CANNOT attend to text tokens (CRITICAL FIX)
+    #     mask[:prefix_len, prefix_len:] = float("-inf")
+        
+    #     # Text tokens can attend to all prefix tokens
+    #     # mask[prefix_len:, :prefix_len] remains 0 (can attend)
+        
+    #     # Text tokens attend to previous text tokens only (causal)
+    #     text_causal_mask = torch.triu(torch.ones(text_len, text_len), diagonal=1) * float("-inf")
+    #     mask[prefix_len:, prefix_len:] = text_causal_mask
+        
+    #     return mask
+
+    def build_prefix_causal_mask(self, prefix_len: int, text_len: int):
+        """Additive mask; 0 = attend, NEG = block (fp32 for stability)."""
+        total_len = prefix_len + text_len
+        # build in fp32; move to GPU later with .to(device=...) but DON'T cast dtype
+        mask = torch.zeros(total_len, total_len, dtype=torch.float32)
+
+        # large finite negative (safer than -inf with fp16/bf16 + fused kernels)
+        NEG = -torch.finfo(mask.dtype).max
+
+        # Prefix → Text: block
+        mask[:prefix_len, prefix_len:] = NEG
+
+        # Text → Text: causal (block future). Use masked_fill, not multiply by -inf.
+        tri = torch.triu(torch.ones(text_len, text_len, dtype=torch.bool), diagonal=1)
+        mask[prefix_len:, prefix_len:].masked_fill_(tri, NEG)
+        return mask
+
+    def forward_intermediates(
+            self,
+            x: torch.Tensor,
+            attn_mask: Optional[torch.Tensor] = None,
+            indices: Optional[Union[int, List[int]]] = None,
+            stop_early: bool = False,
+    ):
+        assert False, "Not currently implemented for ConcatMultimodalTransformer"
+
+    def forward(self, image_embs, text_embs, condition_embs=None):
+        """Forward pass with concatenated embeddings.
+        
+        Args:
+            image_embs: (batch_size, num_image_tokens, width)
+            text_embs: (batch_size, num_text_tokens, width)
+            condition_embs: Optional (batch_size, num_condition_tokens, width)
+                Additional conditioning tokens that will be prepended before image tokens.
+                These tokens receive full bidirectional attention like image tokens.
+            
+        Returns:
+            Text decoder outputs: (batch_size, num_text_tokens, output_dim)
+        """
+        batch_size = text_embs.shape[0]
+        text_len = text_embs.shape[1]
+        
+        # Guard: text length must not exceed context length
+        assert text_len <= self.context_length, \
+            f"text_len ({text_len}) must be <= context_length ({self.context_length})"
+        
+        # Build prefix: [condition_tokens (optional), image_tokens]
+        # All prefix tokens get full bidirectional attention
+        if condition_embs is not None:
+            condition_len = condition_embs.shape[1]
+            
+            # Safety check: condition_len must not exceed the pre-configured condition_prefix_len
+            assert condition_len <= self.condition_prefix_len, \
+                f"condition_len ({condition_len}) exceeds condition_prefix_len ({self.condition_prefix_len})"
+            
+            prefix = torch.cat([condition_embs, image_embs], dim=1)  # (batch, cond_len + img_len, width)
+        else:
+            condition_len = 0
+            prefix = image_embs
+        
+        prefix_len = prefix.shape[1]  # Total prefix length (condition + image tokens)
+        
+        # Concatenate prefix and text embeddings
+        x = torch.cat([prefix, text_embs], dim=1)  # (batch, prefix_len + text_len, width)
+        
+        if not self.batch_first:
+            x = x.permute(1, 0, 2)  # NLD -> LND
+        
+        # Build or retrieve cached prefix-causal attention mask
+        # Dynamically rebuilds when prefix_len changes (handles variable condition_len or image_len)
+        if self._cached_prefix_len != prefix_len or self._cached_attn_mask.numel() == 0:
+            # Build mask for max possible text length (context_length)
+            mask = self.build_prefix_causal_mask(prefix_len, self.context_length)
+            
+            # Directly update the buffer (already registered in __init__)
+            self._cached_attn_mask = mask
+            self._cached_prefix_len = prefix_len
+        
+        # Slice cached mask to actual sequence length
+        seq_len = prefix_len + text_len
+        attn_mask = self._cached_attn_mask[:seq_len, :seq_len].to(device=x.device)
+        
+        # Apply transformer layers with unified self-attention
+        for resblock in self.resblocks:
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                x = checkpoint(resblock, x, None, None, attn_mask, use_reentrant=False)
+            else:
+                x = resblock(x, attn_mask=attn_mask)
+
+        if not self.batch_first:
+            x = x.permute(1, 0, 2)  # LND -> NLD
+
+        # Apply final layer norm
+        x = self.ln_final(x)
+        
+        # Extract only the text portion (skip image prefix)
+        text_output = x[:, prefix_len:, :]  # (batch, text_len, width)
+        
+        # Project to output dimension
+        if self.text_projection is not None:
+            text_output = text_output @ self.text_projection
+        
+        return text_output
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.grad_checkpointing = enable
+
+
+def lock_text_tower(
+    model: nn.Module,
+    unlocked_layers: int = 0,
+):
+    """
+    Lock text tower layers for CLIP models.
+
+    Works with both model architectures:
+    - CustomTextCLIP where text components are in self.text
+    - Standard CLIP where text components are unpacked as attributes
+
+    Args:
+        model: The CLIP model or TextTransformer module
+        unlocked_layers: Number of layers to leave unlocked (from the end)
+    """
+    # Determine where to look for text components
+    if hasattr(model, 'text'):
+        # CustomTextCLIP or already a TextTransformer with nested structure
+        text_module = model.text
+    else:
+        # Standard CLIP or direct TextTransformer
+        text_module = model
+
+    # Collect text components
+    text_params = {}
+    text_params['token_embedding'] = getattr(text_module, 'token_embedding', None)
+    text_params['positional_embedding'] = getattr(text_module, 'positional_embedding', None)
+    text_params['cls_emb'] = getattr(text_module, 'cls_emb', None)
+    text_params['transformer'] = getattr(text_module, 'transformer', None)
+    text_params['ln_final'] = getattr(text_module, 'ln_final', None)
+    text_params['text_projection'] = getattr(text_module, 'text_projection', None)
+
+    # Filter out None values
+    text_params = {k: v for k, v in text_params.items() if v is not None}
+
+    # Freeze all text parameters first
+    for module in text_params.values():
+        if isinstance(module, nn.Parameter):
+            module.requires_grad = False
+        elif isinstance(module, nn.Module):
+            for param in module.parameters():
+                param.requires_grad = False
+
+    if unlocked_layers == 0:
+        return
+
+    # Check if we have transformer blocks to work with
+    transformer = text_params['transformer']
+    if not transformer or not hasattr(transformer, 'resblocks'):
+        return
+
+    total_layers = len(transformer.resblocks)
+    if total_layers == 0:
+        return
+
+    # Build groups for selective unlocking
+    groups = []
+
+    # Group 1: Embeddings
+    embedding_group = []
+    for key in ['token_embedding', 'positional_embedding', 'cls_emb']:
+        if key in text_params:
+            embedding_group.append(text_params[key])
+    if embedding_group:
+        groups.append(embedding_group)
+
+    # Group 2-N: Individual transformer blocks (except last)
+    if total_layers > 1:
+        for block in transformer.resblocks[:-1]:
+            groups.append([block])
+
+    # Combine last transformer block + final ln as the penultimate group
+    last_block = [transformer.resblocks[-1]]
+    if 'ln_final' in text_params:
+        last_block.append(text_params['ln_final'])
+    groups.append(last_block)
+
+    # The final group is the projection only
+    if 'text_projection' in text_params:
+        groups.append([text_params['text_projection']])
+
+    # Helper function to unlock parameters
+    def _unlock(module):
+        if isinstance(module, Sequence):
+            for m in module:
+                _unlock(m)
+        elif isinstance(module, nn.Parameter):
+            module.requires_grad = True
+        elif isinstance(module, nn.Module):
+            for name, param in module.named_parameters():
+                param.requires_grad = True
+
+    # Unlock the specified number of layer groups from the end
+    num_groups_to_unlock = min(unlocked_layers, len(groups))
+    for group in groups[-num_groups_to_unlock:]:
+        _unlock(group)