Update modeling_caption_bert.py

modeling_caption_bert.py

@@ -1,79 +1,167 @@
 # ============================================================================
-# CaptionEncoder: Standalone Consensus-Distilled Caption Embedding Model
-#
-# Produces 768-dim L2-normalized embeddings in geometric consensus space.
-# Trained via distillation from 5-BERT pentachoron consensus.
-# No expert models needed at inference.
+# CaptionBERT-8192: HuggingFace AutoModel-Compatible Implementation
 #
 # Usage:
-#   from caption_encoder import CaptionEncoder
-#   model = CaptionEncoder()
-#   model.load_state_dict(torch.load("best_model.pt"))
-#   # tokenize with bert-base-uncased tokenizer
-#   embedding = model(input_ids, attention_mask)  # (B, 768) L2-normalized
+#   from transformers import AutoModel, AutoTokenizer
+#   model = AutoModel.from_pretrained("AbstractPhil/geolip-captionbert-8192",
+#                                      trust_remote_code=True)
+#   tokenizer = AutoTokenizer.from_pretrained("AbstractPhil/geolip-captionbert-8192")
+#   inputs = tokenizer("A cat on a windowsill", return_tensors="pt",
+#                       padding=True, truncation=True, max_length=512)
+#   outputs = model(**inputs)
+#   embedding = outputs.last_hidden_state  # (B, 768) L2-normalized
 # ============================================================================
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from transformers import PretrainedConfig, PreTrainedModel
+
+
+class CaptionBertConfig(PretrainedConfig):
+    model_type = "caption_bert"
+
+    def __init__(
+        self,
+        vocab_size=30522,
+        max_position_embeddings=8192,
+        hidden_size=384,
+        num_attention_heads=6,
+        num_hidden_layers=6,
+        intermediate_size=1536,
+        output_dim=768,
+        hidden_dropout_prob=0.1,
+        pad_token_id=0,
+        **kwargs,
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.num_hidden_layers = num_hidden_layers
+        self.intermediate_size = intermediate_size
+        self.output_dim = output_dim
+        self.hidden_dropout_prob = hidden_dropout_prob
 
 
-class CaptionEncoder(nn.Module):
+class CaptionBertModel(PreTrainedModel):
     """
-    Standalone transformer caption encoder.
-    No pretrained weights required. Trained via geometric consensus distillation.
+    Consensus-distilled caption encoder.
 
-    The embedding space is the geometric intersection of 5 BERT-family models:
-    BERT-base, ModernBERT-base, RoBERTa-base, ALBERT-base-v2, DistilBERT-base.
-    Aligned via whitened Procrustes rotation. Regularized by pentachoron CV.
+    Produces L2-normalized 768-dim embeddings in the geometric consensus
+    space of 5 BERT-family models (BERT, ModernBERT, RoBERTa, ALBERT, DistilBERT).
 
-    At inference: bert-base-uncased tokenizer + this model.
-    Output: (B, 768) L2-normalized embedding in consensus space.
+    Output:
+        last_hidden_state: (B, output_dim) L2-normalized embedding
+        pooler_output:     (B, output_dim) same as last_hidden_state (for compatibility)
     """
-    def __init__(self, vocab_size=30522, max_len=8192, d_model=384,
-                 n_heads=6, n_layers=6, d_ff=1536, output_dim=768,
-                 dropout=0.1, pad_token_id=0):
-        super().__init__()
-        self.pad_token_id = pad_token_id
-        self.d_model = d_model
-        self.max_len = max_len
-
-        self.token_emb = nn.Embedding(vocab_size, d_model, padding_idx=pad_token_id)
-        self.pos_emb = nn.Embedding(max_len, d_model)
-        self.emb_norm = nn.LayerNorm(d_model)
-        self.emb_drop = nn.Dropout(dropout)
+    config_class = CaptionBertConfig
 
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+
+        # Embeddings
+        self.token_emb = nn.Embedding(
+            config.vocab_size, config.hidden_size,
+            padding_idx=config.pad_token_id)
+        self.pos_emb = nn.Embedding(
+            config.max_position_embeddings, config.hidden_size)
+        self.emb_norm = nn.LayerNorm(config.hidden_size)
+        self.emb_drop = nn.Dropout(config.hidden_dropout_prob)
+
+        # Transformer encoder
         encoder_layer = nn.TransformerEncoderLayer(
-            d_model=d_model, nhead=n_heads, dim_feedforward=d_ff,
-            dropout=dropout, activation="gelu", batch_first=True,
-            norm_first=True)
-        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
+            d_model=config.hidden_size,
+            nhead=config.num_attention_heads,
+            dim_feedforward=config.intermediate_size,
+            dropout=config.hidden_dropout_prob,
+            activation="gelu",
+            batch_first=True,
+            norm_first=True,
+        )
+        self.encoder = nn.TransformerEncoder(
+            encoder_layer, num_layers=config.num_hidden_layers)
 
+        # Output projection to consensus space
         self.output_proj = nn.Sequential(
-            nn.Linear(d_model, d_model),
+            nn.Linear(config.hidden_size, config.hidden_size),
             nn.GELU(),
-            nn.LayerNorm(d_model),
-            nn.Linear(d_model, output_dim),
+            nn.LayerNorm(config.hidden_size),
+            nn.Linear(config.hidden_size, config.output_dim),
         )
 
-    def forward(self, input_ids, attention_mask=None):
+        self.post_init()
+
+    def forward(self, input_ids=None, attention_mask=None, **kwargs):
         B, L = input_ids.shape
-        positions = torch.arange(L, device=input_ids.device).unsqueeze(0)
+        device = input_ids.device
 
+        # Embed
+        positions = torch.arange(L, device=device).unsqueeze(0)
         x = self.token_emb(input_ids) + self.pos_emb(positions)
         x = self.emb_drop(self.emb_norm(x))
 
+        # Transformer with padding mask
         if attention_mask is not None:
-            kpm = ~attention_mask.bool()
+            key_padding_mask = ~attention_mask.bool()
         else:
-            kpm = (input_ids == self.pad_token_id)
+            key_padding_mask = (input_ids == self.config.pad_token_id)
 
-        x = self.encoder(x, src_key_padding_mask=kpm)
+        x = self.encoder(x, src_key_padding_mask=key_padding_mask)
 
+        # Mean pool over non-padding tokens
         if attention_mask is not None:
             mask = attention_mask.unsqueeze(-1).float()
         else:
-            mask = (~kpm).unsqueeze(-1).float()
+            mask = (~key_padding_mask).unsqueeze(-1).float()
         pooled = (x * mask).sum(1) / mask.sum(1).clamp(min=1)
 
-        return F.normalize(self.output_proj(pooled), dim=-1)
+        # Project and normalize
+        embedding = F.normalize(self.output_proj(pooled), dim=-1)
+
+        # Return in HuggingFace-compatible format
+        return type('Output', (), {
+            'last_hidden_state': embedding,
+            'pooler_output': embedding,
+        })()
+
+    def encode(self, texts, tokenizer=None, max_length=512, batch_size=128,
+               device=None):
+        """
+        Convenience method: raw text → L2-normalized embeddings.
+
+        Args:
+            texts: str or list of str
+            tokenizer: AutoTokenizer instance (loads default if None)
+            max_length: max token length
+            batch_size: encoding batch size
+            device: torch device
+
+        Returns:
+            (N, 768) L2-normalized tensor
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+
+        if tokenizer is None:
+            from transformers import AutoTokenizer
+            tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
+
+        if device is None:
+            device = next(self.parameters()).device
+
+        self.eval()
+        all_emb = []
+        with torch.no_grad():
+            for i in range(0, len(texts), batch_size):
+                batch = texts[i:i+batch_size]
+                inputs = tokenizer(
+                    batch, max_length=max_length, padding="max_length",
+                    truncation=True, return_tensors="pt"
+                ).to(device)
+                out = self(input_ids=inputs["input_ids"],
+                          attention_mask=inputs["attention_mask"])
+                all_emb.append(out.last_hidden_state.cpu())
+        return torch.cat(all_emb)