Skip to content

clip.marsclip_text

marsclip_text

Text composition, tokenization, and batching utilities for MarsCLIP.

SimpleTextTokenizer dataclass

SimpleTextTokenizer(vocab: dict[str, int], pad_token: str = '[PAD]', unk_token: str = '[UNK]', bos_token: str = '[BOS]', eos_token: str = '[EOS]')

A lightweight whitespace tokenizer for early MarsCLIP prototyping.

encode

encode(text: str, max_length: int = 64) -> tuple[torch.Tensor, torch.Tensor]

Encode one text string into padded token ids and attention mask.

Source code in src/clip/marsclip_text.py
def encode(self, text: str, max_length: int = 64) -> tuple[torch.Tensor, torch.Tensor]:
    """Encode one text string into padded token ids and attention mask."""
    tokens = self._tokenize(text)
    ids = [self.bos_token_id]
    ids.extend(self.vocab.get(tok, self.unk_token_id) for tok in tokens)
    ids.append(self.eos_token_id)
    ids = ids[:max_length]

    attention = [1] * len(ids)
    if len(ids) < max_length:
        pad_len = max_length - len(ids)
        ids.extend([self.pad_token_id] * pad_len)
        attention.extend([0] * pad_len)

    return (
        torch.tensor(ids, dtype=torch.long),
        torch.tensor(attention, dtype=torch.bool),
    )

batch_encode

batch_encode(texts: list[str], max_length: int = 64) -> tuple[torch.Tensor, torch.Tensor]

Encode a list of texts into batched token ids and attention masks.

Source code in src/clip/marsclip_text.py
def batch_encode(
    self,
    texts: list[str],
    max_length: int = 64,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Encode a list of texts into batched token ids and attention masks."""
    ids, masks = zip(*(self.encode(text, max_length=max_length) for text in texts))
    return torch.stack(list(ids)), torch.stack(list(masks))

MarsCLIPBatchCollator

MarsCLIPBatchCollator(tokenizer: SimpleTextTokenizer, *, text_mode: str = 'raw_plus_expanded', max_length: int = 64)

Batch MarsCLIPDataset samples and tokenize their text fields.

Source code in src/clip/marsclip_text.py
def __init__(
    self,
    tokenizer: SimpleTextTokenizer,
    *,
    text_mode: str = "raw_plus_expanded",
    max_length: int = 64,
) -> None:
    self.tokenizer = tokenizer
    self.text_mode = text_mode
    self.max_length = max_length

compose_rationale_text

compose_rationale_text(rationale_raw: str, rationale_expanded: str | None = None, *, text_mode: str = 'raw_plus_expanded') -> str

Compose the training text from raw and optionally expanded rationale.

Source code in src/clip/marsclip_text.py
def compose_rationale_text(
    rationale_raw: str,
    rationale_expanded: str | None = None,
    *,
    text_mode: str = "raw_plus_expanded",
) -> str:
    """Compose the training text from raw and optionally expanded rationale."""
    raw = str(rationale_raw).strip()
    expanded = None if rationale_expanded is None else str(rationale_expanded).strip()

    if text_mode == "raw":
        return raw
    if text_mode == "expanded":
        return expanded or raw
    if text_mode == "raw_plus_expanded":
        if expanded:
            return f"{raw}\n\nExpanded context: {expanded}"
        return raw
    raise ValueError(f"Unsupported text_mode: {text_mode}")

build_tokenizer_from_manifest

build_tokenizer_from_manifest(manifest: DataFrame, *, text_mode: str = 'raw_plus_expanded', min_freq: int = 1) -> SimpleTextTokenizer

Build a tokenizer from manifest raw/expanded rationale text.

Source code in src/clip/marsclip_text.py
def build_tokenizer_from_manifest(
    manifest: pd.DataFrame,
    *,
    text_mode: str = "raw_plus_expanded",
    min_freq: int = 1,
) -> SimpleTextTokenizer:
    """Build a tokenizer from manifest raw/expanded rationale text."""
    texts = [
        compose_rationale_text(
            row["rationale_desc"],
            row.get("rationale_expanded"),
            text_mode=text_mode,
        )
        for _, row in manifest.iterrows()
    ]
    return SimpleTextTokenizer.build(texts, min_freq=min_freq)