| """ |
| MemoryCLIP — Memory-Extended CLIP-L/14 Text Encoder |
| |
| Single file containing both config and model for HuggingFace AutoModel. |
| No cross-file imports. |
| |
| Usage: |
| from transformers import AutoModel, AutoConfig |
| model = AutoModel.from_pretrained( |
| "AbstractPhil/geolip-clip-vit-large-patch14-ctx576", |
| trust_remote_code=True) |
| emb = model.encode("A long detailed caption...") |
| """ |
|
|
| import math |
| from typing import Optional, List |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import numpy as np |
| from transformers import PretrainedConfig, PreTrainedModel, CLIPTextModel, CLIPTokenizer |
| from transformers.modeling_outputs import BaseModelOutput |
|
|
|
|
| |
| |
| |
|
|
| class MemoryCLIPConfig(PretrainedConfig): |
| model_type = "memory_clip" |
|
|
| def __init__( |
| self, |
| clip_model="openai/clip-vit-large-patch14", |
| clip_hidden=768, |
| clip_layers=12, |
| clip_max_tokens=77, |
| freeze_clip=True, |
| n_memory_tokens=8, |
| bank_size=64, |
| anchor_dim=768, |
| n_bank_heads=8, |
| bank_cross_layers=2, |
| gate_type="gru", |
| extract_layers=(1, 3, 5, 7, 9, 11), |
| layer_fusion="learned", |
| max_content_tokens=18, |
| segment_overlap=4, |
| max_segments=32, |
| cv_target=0.20, |
| **kwargs, |
| ): |
| self.clip_model = clip_model |
| self.clip_hidden = clip_hidden |
| self.clip_layers = clip_layers |
| self.clip_max_tokens = clip_max_tokens |
| self.freeze_clip = freeze_clip |
| self.n_memory_tokens = n_memory_tokens |
| self.bank_size = bank_size |
| self.anchor_dim = anchor_dim |
| self.n_bank_heads = n_bank_heads |
| self.bank_cross_layers = bank_cross_layers |
| self.gate_type = gate_type |
| self.extract_layers = tuple(extract_layers) |
| self.layer_fusion = layer_fusion |
| self.max_content_tokens = max_content_tokens |
| self.segment_overlap = segment_overlap |
| self.max_segments = max_segments |
| self.cv_target = cv_target |
| super().__init__(**kwargs) |
|
|
| @property |
| def n_extract_layers(self): |
| return len(self.extract_layers) |
|
|
| @property |
| def depth_profile_dim(self): |
| return self.n_extract_layers * self.clip_hidden |
|
|
| @property |
| def effective_context(self): |
| return self.max_segments * self.max_content_tokens |
|
|
|
|
| |
| |
| |
|
|
| class GeometricMemoryBank(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.max_size = config.bank_size |
| self.dim = config.anchor_dim |
| self.depth_compressor = nn.Sequential( |
| nn.Linear(config.depth_profile_dim, config.clip_hidden * 2), |
| nn.GELU(), |
| nn.LayerNorm(config.clip_hidden * 2), |
| nn.Linear(config.clip_hidden * 2, config.anchor_dim), |
| ) |
| self.temporal_proj = nn.Linear(1, config.anchor_dim, bias=False) |
| self.cross_attn = nn.ModuleList([ |
| nn.MultiheadAttention(config.clip_hidden, config.n_bank_heads, |
| batch_first=True, dropout=0.1) |
| for _ in range(config.bank_cross_layers) |
| ]) |
| self.cross_norms = nn.ModuleList([ |
| nn.LayerNorm(config.clip_hidden) |
| for _ in range(config.bank_cross_layers) |
| ]) |
| self.cross_ffns = nn.ModuleList([ |
| nn.Sequential( |
| nn.Linear(config.clip_hidden, config.clip_hidden * 2), |
| nn.GELU(), |
| nn.Linear(config.clip_hidden * 2, config.clip_hidden)) |
| for _ in range(config.bank_cross_layers) |
| ]) |
| self.ffn_norms = nn.ModuleList([ |
| nn.LayerNorm(config.clip_hidden) |
| for _ in range(config.bank_cross_layers) |
| ]) |
|
|
| def init_bank(self, batch_size, device): |
| return {"anchors": torch.zeros(batch_size, 0, self.dim, device=device), |
| "n_written": 0} |
|
|
| def write(self, bank, depth_cls, segment_idx=0): |
| B = depth_cls.shape[0] |
| anchor = self.depth_compressor(depth_cls.reshape(B, -1)) |
| anchor = F.normalize(anchor, dim=-1) |
| t = torch.tensor([[segment_idx]], dtype=anchor.dtype, device=anchor.device) |
| anchor = anchor + 0.1 * self.temporal_proj(t / max(self.max_size, 1)) |
| anchor = F.normalize(anchor, dim=-1) |
| anchors = torch.cat([bank["anchors"], anchor.unsqueeze(1)], dim=1) |
| if anchors.shape[1] > self.max_size: |
| anchors = anchors[:, -self.max_size:] |
| return {"anchors": anchors, "n_written": bank["n_written"] + 1, |
| "live_anchor": anchor} |
|
|
| def read(self, memory_tokens, bank): |
| anchors = bank["anchors"] |
| if anchors.shape[1] == 0: |
| return memory_tokens |
| x = memory_tokens |
| for attn, norm, ffn, ffn_norm in zip( |
| self.cross_attn, self.cross_norms, self.cross_ffns, self.ffn_norms): |
| residual = x |
| x, _ = attn(norm(x), anchors, anchors) |
| x = residual + x |
| residual = x |
| x = residual + ffn(ffn_norm(x)) |
| return x |
|
|
|
|
| class DeltaMemoryGate(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| H = config.clip_hidden |
| self.reset_proj = nn.Linear(H * 2, H) |
| self.update_proj = nn.Linear(H * 2, H) |
| self.candidate_proj = nn.Linear(H * 2, H) |
| self.norm = nn.LayerNorm(H) |
|
|
| def forward(self, old, new): |
| cat = torch.cat([old, new], dim=-1) |
| r = torch.sigmoid(self.reset_proj(cat)) |
| z = torch.sigmoid(self.update_proj(cat)) |
| h = torch.tanh(self.candidate_proj(torch.cat([r * old, new], dim=-1))) |
| return self.norm(z * old + (1 - z) * h) |
|
|
|
|
| class LayerFusion(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| n = config.n_extract_layers |
| self.weights = nn.Parameter(torch.ones(n) / n) |
| self.proj = nn.Linear(config.clip_hidden, config.clip_hidden) |
| self.norm = nn.LayerNorm(config.clip_hidden) |
|
|
| def forward(self, layer_outputs): |
| w = F.softmax(self.weights, dim=0) |
| stacked = torch.stack(layer_outputs) |
| fused = (stacked * w.view(-1, 1, 1, 1)).sum(0) |
| return self.norm(self.proj(fused)) |
|
|
|
|
| class TeacherProjector(nn.Module): |
| def __init__(self, student_dim, teacher_dim): |
| super().__init__() |
| self.proj = nn.Linear(student_dim, teacher_dim, bias=True) |
|
|
| def forward(self, x): |
| return self.proj(x) |
|
|
|
|
| |
| |
| |
|
|
| def segment_text(text, clip_tokenizer, max_content=18, overlap=4, max_segments=32): |
| full_tokens = clip_tokenizer.encode(text, add_special_tokens=False) |
| segments = [] |
| stride = max_content - overlap |
| pos = 0 |
| while pos < len(full_tokens) and len(segments) < max_segments: |
| end = min(pos + max_content, len(full_tokens)) |
| chunk = full_tokens[pos:end] |
| sos = clip_tokenizer.bos_token_id or 49406 |
| eos = clip_tokenizer.eos_token_id or 49407 |
| input_ids = [sos] + chunk + [eos] |
| n_pad = 77 - len(input_ids) |
| if n_pad > 0: |
| input_ids = input_ids + [0] * n_pad |
| else: |
| input_ids = input_ids[:77] |
| mask = [1] * min(len(chunk) + 2, 77) + [0] * max(n_pad, 0) |
| mask = mask[:77] |
| segments.append({ |
| "input_ids": torch.tensor(input_ids, dtype=torch.long), |
| "attention_mask": torch.tensor(mask, dtype=torch.long), |
| }) |
| if end >= len(full_tokens): |
| break |
| pos += stride |
| return segments |
|
|
|
|
| |
| |
| |
|
|
| class MemoryCLIPModel(PreTrainedModel): |
| """ |
| Memory-Extended CLIP-L/14 Text Encoder. |
| |
| Extends CLIP's 77-token context to 576 effective tokens via |
| geometric memory bank with depth-profile anchors. |
| """ |
| config_class = MemoryCLIPConfig |
| supports_gradient_checkpointing = False |
|
|
| def __init__(self, config): |
| super().__init__(config) |
|
|
| self.memory_embeddings = nn.Parameter( |
| torch.randn(1, config.n_memory_tokens, config.clip_hidden) * 0.02) |
| self.layer_fusion = LayerFusion(config) |
| self.bank = GeometricMemoryBank(config) |
| self.gate = DeltaMemoryGate(config) |
|
|
| self.output_proj = nn.Sequential( |
| nn.Linear(config.clip_hidden, config.clip_hidden), |
| nn.GELU(), nn.LayerNorm(config.clip_hidden)) |
| self.memory_output_fusion = nn.Sequential( |
| nn.Linear(config.clip_hidden * 2, config.clip_hidden), |
| nn.GELU(), |
| nn.Linear(config.clip_hidden, config.clip_hidden)) |
|
|
| self.clip_cross_attn = nn.ModuleList([ |
| nn.MultiheadAttention(config.clip_hidden, config.n_bank_heads, |
| batch_first=True, dropout=0.1) |
| for _ in range(config.bank_cross_layers) |
| ]) |
| self.clip_cross_norms = nn.ModuleList([ |
| nn.LayerNorm(config.clip_hidden) |
| for _ in range(config.bank_cross_layers) |
| ]) |
| self.clip_cross_ffns = nn.ModuleList([ |
| nn.Sequential( |
| nn.Linear(config.clip_hidden, config.clip_hidden * 2), |
| nn.GELU(), |
| nn.Linear(config.clip_hidden * 2, config.clip_hidden)) |
| for _ in range(config.bank_cross_layers) |
| ]) |
| self.clip_cross_ffn_norms = nn.ModuleList([ |
| nn.LayerNorm(config.clip_hidden) |
| for _ in range(config.bank_cross_layers) |
| ]) |
|
|
| self.proj_modern = TeacherProjector(config.clip_hidden, 1024) |
|
|
| self._clip_text = None |
| self._clip_tokenizer = None |
|
|
| self.post_init() |
|
|
| @property |
| def clip_text(self): |
| if self._clip_text is None: |
| self._clip_text = CLIPTextModel.from_pretrained(self.config.clip_model) |
| self._clip_text.config.output_hidden_states = True |
| for p in self._clip_text.parameters(): |
| p.requires_grad = False |
| device = self.memory_embeddings.device |
| self._clip_text = self._clip_text.to(device) |
| return self._clip_text |
|
|
| @property |
| def clip_tokenizer(self): |
| if self._clip_tokenizer is None: |
| self._clip_tokenizer = CLIPTokenizer.from_pretrained(self.config.clip_model) |
| return self._clip_tokenizer |
|
|
| def init_state(self, batch_size, device=None): |
| if device is None: |
| device = self.memory_embeddings.device |
| return { |
| "memory": self.memory_embeddings.expand(batch_size, -1, -1).clone(), |
| "bank": self.bank.init_bank(batch_size, device), |
| "segment_idx": 0, |
| } |
|
|
| def forward_segment(self, input_ids, attention_mask, state): |
| B = input_ids.shape[0] |
| memory_state = state["memory"] |
| bank = state["bank"] |
| seg_idx = state["segment_idx"] |
|
|
| memory_tokens = self.bank.read(memory_state, bank) |
|
|
| max_len = self.config.clip_max_tokens |
| with torch.no_grad(): |
| clip_out = self.clip_text( |
| input_ids=input_ids[:, :max_len], |
| attention_mask=attention_mask[:, :max_len], |
| output_hidden_states=True, return_dict=True) |
|
|
| all_hiddens = clip_out.hidden_states |
| selected = [all_hiddens[i + 1] for i in self.config.extract_layers] |
| fused = self.layer_fusion(selected) |
|
|
| mem_enriched = memory_tokens |
| for attn, norm, ffn, ffn_norm in zip( |
| self.clip_cross_attn, self.clip_cross_norms, |
| self.clip_cross_ffns, self.clip_cross_ffn_norms): |
| residual = mem_enriched |
| mem_enriched, _ = attn(norm(mem_enriched), fused, fused) |
| mem_enriched = residual + mem_enriched |
| residual = mem_enriched |
| mem_enriched = residual + ffn(ffn_norm(mem_enriched)) |
|
|
| depth_cls = torch.stack([h[:, 1, :] for h in selected], dim=1) |
| new_memory = self.gate(memory_state, mem_enriched) |
| new_bank = self.bank.write(bank, depth_cls, seg_idx) |
|
|
| clip_pooled = clip_out.pooler_output |
| if clip_pooled is None: |
| clip_pooled = clip_out.last_hidden_state[:, -1, :] |
| cls_output = self.output_proj(clip_pooled) |
| memory_delta = self.memory_output_fusion( |
| torch.cat([cls_output, new_memory.mean(dim=1)], dim=-1)) |
| fused_output = cls_output + memory_delta |
|
|
| new_state = { |
| "memory": new_memory, |
| "bank": {"anchors": new_bank["anchors"], |
| "n_written": new_bank["n_written"]}, |
| "segment_idx": seg_idx + 1, |
| } |
| return fused_output, new_state |
|
|
| def forward( |
| self, |
| input_ids: Optional[torch.LongTensor] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| texts: Optional[List[str]] = None, |
| return_dict: bool = True, |
| **kwargs, |
| ) -> BaseModelOutput: |
| """ |
| Accepts either: |
| - input_ids + attention_mask (single 77-token segment) |
| - texts (list of strings, auto-segmented for long context) |
| |
| Returns BaseModelOutput with last_hidden_state = (B, 1, 768). |
| """ |
| device = self.memory_embeddings.device |
|
|
| if texts is not None: |
| embeddings = [self._encode_single(t, device) for t in texts] |
| last_hidden = torch.stack(embeddings) |
| elif input_ids is not None: |
| state = self.init_state(input_ids.shape[0], device) |
| last_hidden, _ = self.forward_segment( |
| input_ids.to(device), attention_mask.to(device), state) |
| else: |
| raise ValueError("Provide either input_ids or texts") |
|
|
| if return_dict: |
| return BaseModelOutput( |
| last_hidden_state=last_hidden.unsqueeze(1), |
| hidden_states=None, attentions=None) |
| return (last_hidden.unsqueeze(1),) |
|
|
| def _encode_single(self, text, device): |
| segments = segment_text( |
| text, self.clip_tokenizer, |
| self.config.max_content_tokens, |
| self.config.segment_overlap, |
| self.config.max_segments) |
| state = self.init_state(1, device) |
| output = None |
| for seg in segments: |
| ids = seg["input_ids"].unsqueeze(0).to(device) |
| mask = seg["attention_mask"].unsqueeze(0).to(device) |
| output, state = self.forward_segment(ids, mask, state) |
| return output.squeeze(0) |
|
|
| def encode(self, texts, batch_size=32, show_progress=False): |
| """ |
| Encode text(s) to 768-dim CLIP-compatible embeddings. |
| |
| Returns: |
| torch.Tensor: (N, 768) or (768,) for single string |
| """ |
| device = self.memory_embeddings.device |
| single = isinstance(texts, str) |
| if single: |
| texts = [texts] |
|
|
| all_embs = [] |
| iterator = range(0, len(texts), batch_size) |
| if show_progress: |
| from tqdm import tqdm |
| iterator = tqdm(iterator, desc="Encoding") |
|
|
| with torch.no_grad(): |
| for i in iterator: |
| batch = texts[i:i + batch_size] |
| embs = [self._encode_single(t, device) for t in batch] |
| all_embs.append(torch.stack(embs)) |
|
|
| result = torch.cat(all_embs, dim=0) |
| return result.squeeze(0) if single else result |
