Upload embedder.py with huggingface_hub

embedder.py

@@ -0,0 +1,79 @@
+from typing import List
+from typing import Tuple
+
+import numpy as np
+import torch
+from torch import nn
+from transformers import AutoTokenizer, T5EncoderModel
+
+
+class TextEmbedder(nn.Module):
+    """
+    Minimal wrapper around a *frozen* T5-base (or any seq-to-seq encoder).
+
+    ────────────  Usage  ─────────────
+    >>> txt = TextEmbedder("google-t5/t5-base", max_len=128)
+    >>> ids, mask = txt.tokenize("some text")
+    >>> z_txt    = txt.encode(torch.tensor(ids)[None, :],
+                              torch.tensor(mask)[None, :])
+    """
+    def __init__(
+        self,
+        model_name: str = "google-t5/t5-base",
+        max_len: int = 128,
+        dtype: torch.dtype = torch.float16,   # keeps weights ≈2× smaller
+    ):
+        super(TextEmbedder, self).__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.encoder   = T5EncoderModel.from_pretrained(model_name).to(dtype=dtype)
+        self.encoder.eval()
+        for p in self.encoder.parameters():   # freeze
+            p.requires_grad = False
+
+        self.max_len = max_len
+
+    # ------------------------------------------------------------ tokenize
+    @torch.no_grad()
+    def tokenize(self, text: str | List[str]) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        padded position are set to 0, and attention mask to 1.
+        For example, if max_len=5 and text="hellow world",
+        the tokenizer will return:
+        input_ids = [21820, 296,  0,  0,  0]
+        attention_mask = [1, 1, 0, 0, 0]
+
+        Returns  (input_ids[int32], attn_mask[int8])  shaped (max_len,)
+        (If *text* is a list, shapes are (B,max_len).)
+        """
+        batch = self.tokenizer(
+            text,
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_len,
+            return_attention_mask=True,
+        )
+        ids   = np.asarray(batch["input_ids"],   dtype=np.int32)
+        mask  = np.asarray(batch["attention_mask"], dtype=np.int8)
+        return ids, mask
+
+    # ------------------------------------------------------------ encode
+    @torch.no_grad()
+    def encode(
+        self,
+        input_ids: torch.LongTensor,        # (B,L)
+        attention_mask: torch.BoolTensor,   # (B,L)
+    ) -> torch.Tensor:                      # (B,L,H)
+        """
+        Run the frozen encoder and return the sequence embeddings
+        (no pooling – do that in your SigLIP loss if you wish).
+        """
+        input_ids     = input_ids
+        attention_mask = attention_mask
+
+        out = self.encoder(input_ids=input_ids,
+                           attention_mask=attention_mask,
+                           return_dict=True)
+        return out.last_hidden_state        # (B,L,hidden)
+
+
+