Spaces:
Sleeping
Sleeping
| """OpenHTR autoregressive decoder with visual cross-attention and KV caching.""" | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| __all__ = ['OpenHTR_Decoder'] | |
| class RMSNorm(nn.Module): | |
| def __init__(self, dim, eps=1e-6): | |
| super().__init__() | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| def forward(self, x): | |
| dtype = x.dtype | |
| x = x.float() | |
| x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) | |
| return self.weight * x.to(dtype) | |
| class SwiGLU(nn.Module): | |
| def __init__(self, dim, hidden_dim, drop=0.0): | |
| super().__init__() | |
| self.w1 = nn.Linear(dim, hidden_dim, bias=False) # gate | |
| self.w3 = nn.Linear(dim, hidden_dim, bias=False) # up | |
| self.w2 = nn.Linear(hidden_dim, dim, bias=False) # down | |
| self.drop = nn.Dropout(drop) | |
| def forward(self, x): | |
| return self.drop(self.w2(F.silu(self.w1(x)) * self.w3(x))) | |
| def build_rope_1d(head_dim, positions, device, dtype=torch.float32, | |
| base=10000.0): | |
| """RoPE cos/sin for the given absolute positions (interleaved pairs). | |
| positions: LongTensor [T] or an int for a single position. | |
| Returns cos/sin of shape [T, head_dim // 2]. | |
| """ | |
| half = head_dim // 2 | |
| if not torch.is_tensor(positions): | |
| positions = torch.tensor([positions], device=device) | |
| inv_freq = 1.0 / (base**(torch.arange(0, head_dim, 2, device=device, | |
| dtype=dtype) / head_dim)) | |
| angles = positions.to(dtype)[:, None] * inv_freq[None, :] # T, half | |
| return angles.cos(), angles.sin() | |
| def rotate_interleaved(x, cos, sin): | |
| """RoPE rotation, interleaved-pair convention. | |
| x: [B, H, T, D]; cos/sin: [T, D//2] (broadcastable). | |
| """ | |
| x1 = x[..., 0::2] | |
| x2 = x[..., 1::2] | |
| out1 = x1 * cos - x2 * sin | |
| out2 = x1 * sin + x2 * cos | |
| return torch.stack([out1, out2], dim=-1).flatten(-2) | |
| class NeoSelfAttention(nn.Module): | |
| """Fused-QKV self-attention with RoPE and KV-cache support.""" | |
| def __init__(self, dim, num_heads, dropout=0.0): | |
| super().__init__() | |
| self.num_heads = num_heads | |
| self.head_dim = dim // num_heads | |
| assert self.head_dim * num_heads == dim | |
| self.qkv = nn.Linear(dim, dim * 3, bias=False) | |
| self.out_proj = nn.Linear(dim, dim, bias=False) | |
| self.dropout = dropout | |
| def forward(self, x, cos=None, sin=None, kv_cache=None, is_causal=False): | |
| B, T, D = x.shape | |
| qkv = self.qkv(x).reshape(B, T, 3, self.num_heads, | |
| self.head_dim).permute(2, 0, 3, 1, 4) | |
| q, k, v = qkv.unbind(0) | |
| if cos is not None: | |
| cos = cos.to(q.dtype)[None, None] # 1 1 T hd/2 | |
| sin = sin.to(q.dtype)[None, None] | |
| q = rotate_interleaved(q, cos, sin) | |
| k = rotate_interleaved(k, cos, sin) | |
| if kv_cache is not None and kv_cache[0] is not None: | |
| k = torch.cat([kv_cache[0], k], dim=2) | |
| v = torch.cat([kv_cache[1], v], dim=2) | |
| y = F.scaled_dot_product_attention( | |
| q, | |
| k, | |
| v, | |
| dropout_p=self.dropout if self.training else 0.0, | |
| is_causal=is_causal, | |
| ) | |
| y = y.transpose(1, 2).reshape(B, T, D) | |
| return self.out_proj(y), k, v | |
| class NeoCrossAttention(nn.Module): | |
| """Cross-attention against a fixed memory (K/V precomputed once).""" | |
| def __init__(self, dim, num_heads, dropout=0.0): | |
| super().__init__() | |
| self.num_heads = num_heads | |
| self.head_dim = dim // num_heads | |
| self.embed_dim = dim | |
| self.q = nn.Linear(dim, dim, bias=False) | |
| self.kv = nn.Linear(dim, dim * 2, bias=False) | |
| self.out_proj = nn.Linear(dim, dim, bias=False) | |
| self.dropout = dropout | |
| def project_kv(self, memory): | |
| """memory: [B, N, D] -> (k, v) each [B, H, N, head_dim].""" | |
| B, N, _ = memory.shape | |
| kv = self.kv(memory).reshape(B, N, 2, self.num_heads, | |
| self.head_dim).permute(2, 0, 3, 1, 4) | |
| return kv.unbind(0) | |
| def forward(self, x, mem_kv): | |
| B, T, D = x.shape | |
| q = self.q(x).reshape(B, T, self.num_heads, | |
| self.head_dim).transpose(1, 2) | |
| k, v = mem_kv | |
| y = F.scaled_dot_product_attention( | |
| q, | |
| k, | |
| v, | |
| dropout_p=self.dropout if self.training else 0.0, | |
| ) | |
| y = y.transpose(1, 2).reshape(B, T, D) | |
| return self.out_proj(y) | |
| class NeoDecoderBlock(nn.Module): | |
| """Pre-norm block: RoPE self-attn -> cross-attn -> SwiGLU FFN.""" | |
| def __init__(self, | |
| d_model, | |
| nhead, | |
| mlp_ratio=2.67, | |
| attention_dropout_rate=0.0, | |
| residual_dropout_rate=0.1, | |
| rope=True): | |
| super().__init__() | |
| hidden = max(64, int(d_model * mlp_ratio) // 64 * 64) | |
| self.self_attn = NeoSelfAttention(d_model, | |
| nhead, | |
| dropout=attention_dropout_rate) | |
| self.norm1 = RMSNorm(d_model) | |
| self.cross_attn = NeoCrossAttention(d_model, | |
| nhead, | |
| dropout=attention_dropout_rate) | |
| self.norm2 = RMSNorm(d_model) | |
| self.ffn = SwiGLU(d_model, hidden, drop=residual_dropout_rate) | |
| self.norm3 = RMSNorm(d_model) | |
| self.dropout = nn.Dropout(residual_dropout_rate) | |
| self.rope = rope | |
| def _rope(self, positions, device, dtype): | |
| if not self.rope: | |
| return None, None | |
| return build_rope_1d(self.self_attn.head_dim, positions, device, | |
| dtype) | |
| def forward_train(self, x, mem_kv): | |
| """Full-sequence causal pass. x: [B, T, D].""" | |
| T = x.shape[1] | |
| pos = torch.arange(T, device=x.device) | |
| cos, sin = self._rope(pos, x.device, x.dtype) | |
| a, _, _ = self.self_attn(self.norm1(x), cos, sin, is_causal=T > 1) | |
| x = x + self.dropout(a) | |
| x = x + self.dropout(self.cross_attn(self.norm2(x), mem_kv)) | |
| x = x + self.dropout(self.ffn(self.norm3(x))) | |
| return x | |
| def step(self, x, mem_kv, kv_cache, pos): | |
| """Single-token incremental pass. | |
| x: [B, 1, D]; kv_cache: [k, v] list updated in place; | |
| pos: absolute RoPE position of the incoming token. | |
| """ | |
| cos, sin = self._rope(pos, x.device, x.dtype) | |
| a, k, v = self.self_attn(self.norm1(x), cos, sin, kv_cache=kv_cache) | |
| kv_cache[0], kv_cache[1] = k, v | |
| x = x + self.dropout(a) | |
| x = x + self.dropout(self.cross_attn(self.norm2(x), mem_kv)) | |
| x = x + self.dropout(self.ffn(self.norm3(x))) | |
| return x | |
| class NeoEncoderBlock(nn.Module): | |
| """Pre-norm RoPE self-attention + SwiGLU block (optional memory side).""" | |
| def __init__(self, | |
| d_model, | |
| nhead, | |
| mlp_ratio=2.67, | |
| attention_dropout_rate=0.0, | |
| residual_dropout_rate=0.1): | |
| super().__init__() | |
| hidden = max(64, int(d_model * mlp_ratio) // 64 * 64) | |
| self.self_attn = NeoSelfAttention(d_model, | |
| nhead, | |
| dropout=attention_dropout_rate) | |
| self.norm1 = RMSNorm(d_model) | |
| self.ffn = SwiGLU(d_model, hidden, drop=residual_dropout_rate) | |
| self.norm2 = RMSNorm(d_model) | |
| self.dropout = nn.Dropout(residual_dropout_rate) | |
| def forward(self, x): | |
| T = x.shape[1] | |
| pos = torch.arange(T, device=x.device) | |
| cos, sin = build_rope_1d(self.self_attn.head_dim, pos, x.device, | |
| x.dtype) | |
| a, _, _ = self.self_attn(self.norm1(x), cos, sin, is_causal=False) | |
| x = x + self.dropout(a) | |
| x = x + self.dropout(self.ffn(self.norm2(x))) | |
| return x | |
| class OpenHTR_Decoder(nn.Module): | |
| """Decode text from visual tokens with causal self-attention and cross-attention.""" | |
| def __init__( | |
| self, | |
| in_channels, | |
| out_channels, | |
| nhead=None, | |
| num_encoder_layers=0, | |
| beam_size=0, # unused, greedy decoding only; kept for config compat | |
| num_decoder_layers=2, | |
| max_len=25, | |
| attention_dropout_rate=0.0, | |
| residual_dropout_rate=0.1, | |
| mlp_ratio=2.67, | |
| **kwargs, | |
| ): | |
| super().__init__() | |
| self.out_channels = out_channels | |
| self.ignore_index = out_channels - 1 | |
| self.bos = out_channels - 2 | |
| self.eos = 0 | |
| self.max_len = max_len | |
| d_model = in_channels | |
| nhead = nhead if nhead is not None else d_model // 64 | |
| assert d_model % nhead == 0, \ | |
| f'd_model ({d_model}) must be divisible by nhead ({nhead})' | |
| self.embedding = nn.Embedding(out_channels, d_model, padding_idx=0) | |
| if num_encoder_layers is not None and num_encoder_layers > 0: | |
| self.encoder = nn.ModuleList([ | |
| NeoEncoderBlock( | |
| d_model, | |
| nhead, | |
| mlp_ratio=mlp_ratio, | |
| attention_dropout_rate=attention_dropout_rate, | |
| residual_dropout_rate=residual_dropout_rate, | |
| ) for _ in range(num_encoder_layers) | |
| ]) | |
| else: | |
| self.encoder = None | |
| self.decoder = nn.ModuleList([ | |
| NeoDecoderBlock( | |
| d_model, | |
| nhead, | |
| mlp_ratio=mlp_ratio, | |
| attention_dropout_rate=attention_dropout_rate, | |
| residual_dropout_rate=residual_dropout_rate, | |
| ) for _ in range(num_decoder_layers) | |
| ]) | |
| self.norm_out = RMSNorm(d_model) | |
| self.tgt_word_prj = nn.Linear(d_model, self.out_channels - 2, | |
| bias=False) | |
| self.d_model = d_model | |
| self.apply(self._init_weights) | |
| nn.init.zeros_(self.embedding.weight.data[0]) # padding idx | |
| def _init_weights(self, m): | |
| if isinstance(m, nn.Linear): | |
| nn.init.normal_(m.weight, mean=0.0, std=0.02) | |
| elif isinstance(m, nn.Embedding): | |
| nn.init.normal_(m.weight, mean=0.0, std=0.02) | |
| def _encode(self, src): | |
| if self.encoder is not None: | |
| for layer in self.encoder: | |
| src = layer(src) | |
| return src | |
| def forward_train(self, src, tgt): | |
| tgt = tgt[:, :-1] | |
| memory = self._encode(src) | |
| mem_kv = [layer.cross_attn.project_kv(memory) | |
| for layer in self.decoder] | |
| x = self.embedding(tgt) | |
| for i, layer in enumerate(self.decoder): | |
| x = layer.forward_train(x, mem_kv[i]) | |
| x = self.norm_out(x) | |
| return self.tgt_word_prj(x) | |
| def forward_test(self, src): | |
| bs = src.shape[0] | |
| device = src.device | |
| memory = self._encode(src) | |
| mem_kv = [layer.cross_attn.project_kv(memory) | |
| for layer in self.decoder] | |
| dec_seq = torch.full((bs, self.max_len + 1), | |
| self.ignore_index, | |
| dtype=torch.int64, | |
| device=device) | |
| dec_seq[:, 0] = self.bos | |
| logits = [] | |
| kv_caches = [[None, None] for _ in self.decoder] | |
| with torch.no_grad(): | |
| for step in range(0, self.max_len): | |
| tok = dec_seq[:, step:step + 1] # B 1 | |
| x = self.embedding(tok) | |
| for i, layer in enumerate(self.decoder): | |
| x = layer.step(x, mem_kv[i], kv_caches[i], pos=step) | |
| x = self.norm_out(x) | |
| word_prob = F.softmax(self.tgt_word_prj(x[:, -1, :]), | |
| dim=-1).unsqueeze(1) # B 1 V | |
| logits.append(word_prob) | |
| dec_seq[:, step + 1] = word_prob.squeeze(1).argmax(-1) | |
| # stop once every sequence has emitted at least one EOS | |
| if (dec_seq == self.eos).any(dim=-1).all(): | |
| break | |
| logits = torch.cat(logits, dim=1) | |
| return logits | |
| def forward(self, src, data=None): | |
| if self.training: | |
| max_len = data[1].max() | |
| tgt = data[0][:, :2 + max_len] | |
| return self.forward_train(src, tgt) | |
| return self.forward_test(src) | |