Replace broken LRSCLIPTextEncoder with official DGTRS architecture

Root cause of NaN: our open_clip wrapper had 3 bugs:
1. Positional embeddings summed for all positions instead of masked
   (official: mask1 for pos 0-19, mask2 for pos 20-247)
2. open_clip uses batch-first transformer, DGTRS uses sequence-first
   (LND format with nn.MultiheadAttention)
3. open_clip tokenizer truncates to 77 tokens, DGTRS needs 248

Fix: copied official DGTRS text encoder architecture from
github.com/MitsuiChen14/DGTRS (Apache-2.0):
- src/models/dgtrs/model.py: DGTRSTextEncoder, build_model,
  load_dgtrs_text_encoder, tokenize_dgtrs
- src/models/dgtrs/simple_tokenizer.py: BPE tokenizer (248 tokens)
- src/models/dgtrs/bpe_simple_vocab_16e6.txt.gz: vocabulary

Removed: LRSCLIPTextEncoder class, open_clip dependency for text encoding

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

src/models/asymmetric_encoder.py

@@ -17,7 +17,6 @@ import warnings
 from pathlib import Path
 
 import coloredlogs
-import open_clip
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -26,6 +25,7 @@ LOGGER = logging.getLogger("caption_test.model")
 coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
 from safetensors.torch import load_file as load_safetensors
 
+from src.models.dgtrs.model import DGTRSTextEncoder, load_dgtrs_text_encoder, tokenize_dgtrs
 from src.models.dual_encoder import GatedFusion, ProjectionHead
 
 
@@ -197,116 +197,8 @@ class DINOv3ViT(nn.Module):
         return model
 
 
-# ---------------------------------------------------------------------------
+# LRSCLIPTextEncoder removed — replaced by official DGTRS architecture
-# LRSCLIP (DGTRS-CLIP) text encoder — open_clip with KPS positional embedding
+# in src/models/dgtrs/model.py (DGTRSTextEncoder)
-# ---------------------------------------------------------------------------
-
-class LRSCLIPTextEncoder(nn.Module):
-    """DGTRS-CLIP text encoder with KPS positional embedding (248 tokens).
-
-    Wraps open_clip ViT-L-14 text tower. Handles the extra
-    positional_embedding_res (KPS residual) not present in vanilla open_clip.
-
-    Checkpoint keys (text-only):
-        token_embedding.weight: [49408, 768]
-        positional_embedding: [248, 768]
-        positional_embedding_res: [248, 768]
-        transformer.resblocks.{0-11}.{attn,mlp,ln}.*
-        ln_final.{weight,bias}
-        text_projection: [768, 768]
-    """
-
-    def __init__(self, context_length: int = 248, embed_dim: int = 768) -> None:
-        super().__init__()
-        self.context_length = context_length
-        self.embed_dim = embed_dim
-
-        # Build the open_clip model to get correct architecture, then
-        # replace positional embedding and add KPS residual.
-        # Suppress "no pretrained weights" warning — we load DGTRS-CLIP weights after.
-        _root = logging.getLogger()
-        _prev_level = _root.level
-        _root.setLevel(logging.ERROR)
-        clip_model = open_clip.create_model("ViT-L-14", pretrained=None)
-        _root.setLevel(_prev_level)
-
-        # Extract text components.
-        self.token_embedding = clip_model.token_embedding
-        self.transformer = clip_model.transformer
-        self.ln_final = clip_model.ln_final
-        self.text_projection = clip_model.text_projection
-        self.attn_mask = None  # rebuilt in forward
-
-        # Replace positional embedding with 248-length version.
-        self.positional_embedding = nn.Parameter(
-            torch.zeros(context_length, embed_dim),
-        )
-        self.positional_embedding_res = nn.Parameter(
-            torch.zeros(context_length, embed_dim),
-        )
-
-    def _build_attn_mask(self, context_length: int, device: torch.device) -> torch.Tensor:
-        mask = torch.empty(context_length, context_length, device=device)
-        mask.fill_(float("-inf"))
-        mask.triu_(1)
-        return mask
-
-    def forward(self, text: torch.Tensor) -> torch.Tensor:
-        """Encode tokenized text.
-
-        Args:
-            text: Token IDs [B, T] (T <= 248).
-
-        Returns:
-            Text embeddings [B, 768], L2-normalized.
-        """
-        T = text.shape[1]
-        x = self.token_embedding(text)  # [B, T, 768]
-        pos = self.positional_embedding[:T] + self.positional_embedding_res[:T]
-        x = x + pos
-
-        if self.attn_mask is None or self.attn_mask.shape[0] != T:
-            self.attn_mask = self._build_attn_mask(T, x.device)
-        attn_mask = self.attn_mask[:T, :T]
-
-        # open_clip transformer expects batch-first [B, T, D].
-        x = self.transformer(x, attn_mask=attn_mask)
-        x = self.ln_final(x)  # [B, T, D]
-
-        # Take features at EOS token position.
-        eos_idx = text.argmax(dim=-1)
-        x = x[torch.arange(x.shape[0], device=x.device), eos_idx]
-
-        # Project.
-        if isinstance(self.text_projection, nn.Parameter):
-            x = x @ self.text_projection
-        else:
-            x = self.text_projection(x)
-
-        return F.normalize(x, dim=-1)
-
-    @classmethod
-    def from_pretrained(cls, path: str | Path) -> LRSCLIPTextEncoder:
-        """Load DGTRS-CLIP checkpoint, extracting only text encoder weights."""
-        model = cls()
-        path = Path(path)
-        LOGGER.info("📝 Loading LRSCLIP text encoder from %s", path.name)
-        full_state = torch.load(str(path), map_location="cpu", weights_only=False)
-        if "state_dict" in full_state:
-            full_state = full_state["state_dict"]
-
-        # Filter out visual.* keys — keep only text encoder.
-        text_state = {
-            k: v for k, v in full_state.items()
-            if not k.startswith("visual.")
-        }
-
-        # Handle text_projection shape: checkpoint may be [768, 768] Parameter
-        # while open_clip stores it as nn.Parameter directly.
-        model.load_state_dict(text_state, strict=False)
-        n_params = sum(p.numel() for p in model.parameters())
-        LOGGER.info("📝 LRSCLIP loaded: %s params, context=%d tokens", f"{n_params:,}", model.context_length)
-        return model
 
 
 # ---------------------------------------------------------------------------
@@ -396,15 +288,13 @@ class AsymmetricEncoder(nn.Module):
         self._freeze(self.drone_encoder)
         self._freeze(self.sat_encoder)
 
-        # Text encoder (partial unfreeze).
+        # Text encoder — official DGTRS architecture (partial unfreeze).
         if not baseline_mode:
-            self.text_encoder = LRSCLIPTextEncoder.from_pretrained(lrsclip_path)
+            self.text_encoder = load_dgtrs_text_encoder(lrsclip_path)
             self._freeze(self.text_encoder)
             self._unfreeze_text_last_block()
-            self.tokenizer = open_clip.get_tokenizer("ViT-L-14")
         else:
             self.text_encoder = None
-            self.tokenizer = None
 
         # Projection heads.
         self.proj_drone = ProjectionHead(
@@ -473,12 +363,8 @@ class AsymmetricEncoder(nn.Module):
         return F.normalize(fused, dim=-1)
 
     def _encode_single_text(self, texts: list[str]) -> torch.Tensor:
-        """Tokenize and encode a list of strings."""
+        """Tokenize and encode a list of strings using DGTRS tokenizer."""
-        tokens = self.tokenizer(list(texts)).to(self.device)
+        tokens = tokenize_dgtrs(list(texts)).to(self.device)
-        # Pad/truncate to context_length.
-        T = tokens.shape[1]
-        if T > self.text_encoder.context_length:
-            tokens = tokens[:, :self.text_encoder.context_length]
         return self.text_encoder(tokens)
 
     def forward(

src/models/dgtrs/__init__.py

@@ -0,0 +1,7 @@
+"""DGTRS-CLIP (LRSCLIP) text encoder.
+
+Adapted from https://github.com/MitsuiChen14/DGTRS (Apache-2.0).
+Only the text encoder components are used.
+"""
+
+from src.models.dgtrs.model import build_model, load_dgtrs_text_encoder, tokenize_dgtrs

src/models/dgtrs/model.py

@@ -0,0 +1,277 @@
+"""DGTRS-CLIP model — text encoder components.
+
+Adapted from https://github.com/MitsuiChen14/DGTRS (Apache-2.0).
+Only text-encoder classes are kept; vision encoder removed.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections import OrderedDict
+from pathlib import Path
+from typing import Union
+
+import coloredlogs
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from src.models.dgtrs.simple_tokenizer import SimpleTokenizer as _Tokenizer
+
+LOGGER = logging.getLogger("caption_test.dgtrs")
+coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
+
+_tokenizer = _Tokenizer()
+
+
+# ---------------------------------------------------------------------------
+# Transformer blocks (DGTRS original — sequence-first LND format)
+# ---------------------------------------------------------------------------
+
+class LayerNorm(nn.LayerNorm):
+    """LayerNorm that handles fp16 by casting to fp32 internally."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model)),
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = (
+            self.attn_mask.to(dtype=x.dtype, device=x.device)
+            if self.attn_mask is not None else None
+        )
+        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.attention(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(
+            *[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)]
+        )
+
+    def forward(self, x: torch.Tensor):
+        return self.resblocks(x)
+
+
+# ---------------------------------------------------------------------------
+# DGTRS-CLIP text encoder
+# ---------------------------------------------------------------------------
+
+class DGTRSTextEncoder(nn.Module):
+    """DGTRS-CLIP text encoder with KPS dual positional embeddings.
+
+    Context length: 248 tokens.
+    Uses mask1 (positions 0-19) for positional_embedding and
+    mask2 (positions 20-247) for positional_embedding_res.
+
+    This is the official architecture from
+    https://github.com/MitsuiChen14/DGTRS/blob/main/model/model_longclip.py
+    """
+
+    CONTEXT_LENGTH = 248
+
+    def __init__(
+        self,
+        vocab_size: int = 49408,
+        transformer_width: int = 768,
+        transformer_heads: int = 12,
+        transformer_layers: int = 12,
+        embed_dim: int = 768,
+    ) -> None:
+        super().__init__()
+        self.context_length = self.CONTEXT_LENGTH
+        self.embed_dim = embed_dim
+
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.positional_embedding = nn.Parameter(
+            torch.empty(self.CONTEXT_LENGTH, transformer_width),
+        )
+        self.positional_embedding_res = nn.Parameter(
+            torch.empty(self.CONTEXT_LENGTH, transformer_width),
+        )
+
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self._build_attention_mask(),
+        )
+
+        self.ln_final = LayerNorm(transformer_width)
+        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+
+        # KPS masks: positions 0-19 use positional_embedding,
+        # positions 20-247 use positional_embedding_res.
+        mask1 = torch.zeros(self.CONTEXT_LENGTH, 1)
+        mask1[:20, :] = 1
+        mask2 = torch.zeros(self.CONTEXT_LENGTH, 1)
+        mask2[20:, :] = 1
+        self.register_buffer("mask1", mask1)
+        self.register_buffer("mask2", mask2)
+
+    def _build_attention_mask(self) -> torch.Tensor:
+        mask = torch.empty(self.CONTEXT_LENGTH, self.CONTEXT_LENGTH)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)
+        return mask
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.token_embedding.weight.dtype
+
+    def forward(self, text: torch.Tensor) -> torch.Tensor:
+        """Encode tokenized text.
+
+        Args:
+            text: Token IDs [B, 248] (int/long).
+
+        Returns:
+            Text embeddings [B, embed_dim].
+        """
+        x = self.token_embedding(text).type(self.dtype)  # [B, 248, D]
+
+        # Dual masked positional embeddings (KPS).
+        x = x + (
+            self.positional_embedding * self.mask1
+            + self.positional_embedding_res * self.mask2
+        ).type(self.dtype)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND (sequence-first for nn.MultiheadAttention)
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        # Take features from EOT token (highest token ID in each sequence).
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+        return x
+
+
+# ---------------------------------------------------------------------------
+# Loading and tokenization
+# ---------------------------------------------------------------------------
+
+def _extract_text_state(full_state: dict) -> dict:
+    """Extract text encoder keys from a full DGTRS-CLIP state dict."""
+    text_keys = {
+        "token_embedding", "positional_embedding", "positional_embedding_res",
+        "transformer", "ln_final", "text_projection",
+    }
+    return {
+        k: v for k, v in full_state.items()
+        if any(k == prefix or k.startswith(prefix + ".") for prefix in text_keys)
+    }
+
+
+def build_model(state_dict: dict) -> DGTRSTextEncoder:
+    """Build DGTRSTextEncoder from a DGTRS-CLIP state dict.
+
+    Auto-detects architecture dimensions from the state dict.
+    """
+    embed_dim = state_dict["text_projection"].shape[1]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(
+        k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks")
+    ))
+
+    model = DGTRSTextEncoder(
+        vocab_size=vocab_size,
+        transformer_width=transformer_width,
+        transformer_heads=transformer_heads,
+        transformer_layers=transformer_layers,
+        embed_dim=embed_dim,
+    )
+    # mask1/mask2 are buffers created in __init__, not in checkpoint.
+    model.load_state_dict(state_dict, strict=False)
+    return model.eval()
+
+
+def load_dgtrs_text_encoder(
+    checkpoint_path: str | Path,
+    device: str = "cpu",
+) -> DGTRSTextEncoder:
+    """Load DGTRS-CLIP text encoder from checkpoint.
+
+    Extracts only text encoder weights from the full CLIP checkpoint.
+    """
+    path = Path(checkpoint_path)
+    LOGGER.info("📝 Loading DGTRS-CLIP text encoder from %s", path.name)
+    full_state = torch.load(str(path), map_location="cpu", weights_only=False)
+    if "state_dict" in full_state:
+        full_state = full_state["state_dict"]
+
+    text_state = _extract_text_state(full_state)
+    model = build_model(text_state)
+    model = model.to(device)
+
+    n_params = sum(p.numel() for p in model.parameters())
+    LOGGER.info(
+        "📝 DGTRS text encoder loaded: %s params, context=%d tokens",
+        f"{n_params:,}", model.context_length,
+    )
+    return model
+
+
+def tokenize_dgtrs(
+    texts: str | list[str],
+    context_length: int = 248,
+    truncate: bool = True,
+) -> torch.Tensor:
+    """Tokenize text for DGTRS-CLIP (248 token context).
+
+    Args:
+        texts: Input string or list of strings.
+        context_length: Output sequence length (default 248).
+        truncate: Whether to truncate long sequences.
+
+    Returns:
+        Token IDs [B, context_length] (int tensor).
+    """
+    if isinstance(texts, str):
+        texts = [texts]
+
+    sot_token = _tokenizer.encoder["<|startoftext|>"]
+    eot_token = _tokenizer.encoder["<|endoftext|>"]
+    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
+
+    result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
+    for i, tokens in enumerate(all_tokens):
+        if len(tokens) > context_length:
+            if truncate:
+                tokens = tokens[:context_length]
+                tokens[-1] = eot_token
+            else:
+                raise RuntimeError(f"Input too long ({len(tokens)} > {context_length})")
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result

src/models/dgtrs/simple_tokenizer.py

@@ -0,0 +1,125 @@
+"""BPE tokenizer for DGTRS-CLIP / LongCLIP.
+
+Adapted from https://github.com/MitsuiChen14/DGTRS (Apache-2.0).
+"""
+
+import gzip
+import html
+import os
+from functools import lru_cache
+
+import ftfy
+import regex as re
+
+
+@lru_cache()
+def default_bpe():
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
+
+
+@lru_cache()
+def bytes_to_unicode():
+    bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8 + n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r"\s+", " ", text)
+    text = text.strip()
+    return text
+
+
+class SimpleTokenizer:
+    def __init__(self, bpe_path: str = default_bpe()):
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        merges = gzip.open(bpe_path).read().decode("utf-8").split("\n")
+        merges = merges[1:49152 - 256 - 2 + 1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(bytes_to_unicode().values())
+        vocab = vocab + [v + "</w>" for v in vocab]
+        for merge in merges:
+            vocab.append("".join(merge))
+        vocab.extend(["<|startoftext|>", "<|endoftext|>"])
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {"<|startoftext|>": "<|startoftext|>", "<|endoftext|>": "<|endoftext|>"}
+        self.pat = re.compile(
+            r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
+            re.IGNORECASE,
+        )
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + (token[-1] + "</w>",)
+        pairs = get_pairs(word)
+        if not pairs:
+            return token + "</w>"
+        while True:
+            bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except ValueError:
+                    new_word.extend(word[i:])
+                    break
+                if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
+                    new_word.append(first + second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = " ".join(word)
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        bpe_tokens = []
+        text = whitespace_clean(basic_clean(text)).lower()
+        for token in re.findall(self.pat, text):
+            token = "".join(self.byte_encoder[b] for b in token.encode("utf-8"))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(" "))
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = "".join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors="replace").replace("</w>", " ")
+        return text