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)
|