| """ |
| PyTorch Vision Encoder module for LawSLM Multimodal Vision-Language Architecture. |
| Converts 3-channel image inputs [B, 3, H, W] into token projection sequence [B, N_patches, d_model]. |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from typing import Tuple, Optional, Dict, Any |
|
|
| from slm.utils.logger import get_logger |
|
|
| logger = get_logger("slm.vision.encoder") |
|
|
|
|
| class VisionPatchEmbedding(nn.Module): |
| """ |
| Splits image tensor into non-overlapping patches and projects them to target d_model dimension. |
| """ |
|
|
| def __init__( |
| self, |
| image_size: int = 224, |
| patch_size: int = 16, |
| in_channels: int = 3, |
| embed_dim: int = 128 |
| ) -> None: |
| super().__init__() |
| self.image_size = image_size |
| self.patch_size = patch_size |
| self.n_patches = (image_size // patch_size) ** 2 |
| self.embed_dim = embed_dim |
|
|
| |
| self.proj = nn.Conv2d( |
| in_channels, |
| embed_dim, |
| kernel_size=patch_size, |
| stride=patch_size |
| ) |
| self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) |
| self.pos_embedding = nn.Parameter(torch.randn(1, self.n_patches + 1, embed_dim) * 0.02) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """ |
| Args: |
| x: Input image tensor of shape [B, 3, H, W] |
| |
| Returns: |
| Sequence embedding tensor of shape [B, N_patches + 1, d_model] |
| """ |
| B, C, H, W = x.shape |
| if H != self.image_size or W != self.image_size: |
| x = F.interpolate(x, size=(self.image_size, self.image_size), mode="bilinear", align_corners=False) |
|
|
| x = self.proj(x) |
| x = x.flatten(2).transpose(1, 2) |
|
|
| cls_tokens = self.cls_token.expand(B, -1, -1) |
| x = torch.cat((cls_tokens, x), dim=1) |
| x = x + self.pos_embedding |
| return x |
|
|
|
|
| class VisionEncoder(nn.Module): |
| """ |
| Lightweight Vision Encoder Transformer Block for image feature extraction. |
| """ |
|
|
| def __init__( |
| self, |
| image_size: int = 224, |
| patch_size: int = 16, |
| in_channels: int = 3, |
| d_model: int = 128, |
| n_layers: int = 2, |
| n_heads: int = 4 |
| ) -> None: |
| super().__init__() |
| self.patch_embed = VisionPatchEmbedding( |
| image_size=image_size, |
| patch_size=patch_size, |
| in_channels=in_channels, |
| embed_dim=d_model |
| ) |
|
|
| encoder_layer = nn.TransformerEncoderLayer( |
| d_model=d_model, |
| nhead=n_heads, |
| dim_feedforward=d_model * 4, |
| dropout=0.1, |
| activation="gelu", |
| batch_first=True |
| ) |
| self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers) |
| self.ln_post = nn.LayerNorm(d_model) |
|
|
| def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: |
| """ |
| Processes normalized image pixels into sequence embeddings. |
| |
| Args: |
| pixel_values: Image tensor [B, 3, H, W] |
| |
| Returns: |
| Image token sequence tensor [B, N_patches + 1, d_model] |
| """ |
| tokens = self.patch_embed(pixel_values) |
| out = self.transformer(tokens) |
| return self.ln_post(out) |
|
|
