Upload src/models/sequence_analyzer.py with huggingface_hub

src/models/sequence_analyzer.py

@@ -0,0 +1,147 @@
+"""
+Sequence Analyzer — Bidirectional LSTM with Self-Attention.
+
+Consumes a sequence of CNN feature embeddings and produces a single
+temporal pattern encoding that captures heat-pattern evolution.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SelfAttention(nn.Module):
+    """
+    Additive (Bahdanau-style) self-attention over a sequence of hidden states.
+
+    Learns which timesteps are most informative and produces a
+    weighted context vector.
+    """
+
+    def __init__(self, hidden_size: int):
+        super().__init__()
+        self.attention_fc = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size // 2),
+            nn.Tanh(),
+            nn.Linear(hidden_size // 2, 1, bias=False),
+        )
+
+    def forward(
+        self, hidden_states: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            hidden_states: (B, T, H)
+
+        Returns:
+            context:  (B, H)  — weighted sum
+            weights:  (B, T)  — attention weights (for visualisation)
+        """
+        scores = self.attention_fc(hidden_states).squeeze(-1)  # (B, T)
+        weights = F.softmax(scores, dim=1)                     # (B, T)
+        context = torch.bmm(
+            weights.unsqueeze(1), hidden_states
+        ).squeeze(1)  # (B, H)
+        return context, weights
+
+
+class SequenceAnalyzer(nn.Module):
+    """
+    Bidirectional LSTM + Self-Attention for temporal analysis
+    of CNN feature sequences.
+
+    Architecture:
+        Input features (B, T, D)
+          → LayerNorm
+          → Bi-LSTM (2 layers, hidden=128)
+          → Self-Attention → context (B, 2*hidden)
+          → FC projection → (B, output_dim=256)
+    """
+
+    def __init__(
+        self,
+        input_dim: int = 256,
+        hidden_size: int = 128,
+        num_layers: int = 2,
+        output_dim: int = 256,
+        bidirectional: bool = True,
+        dropout: float = 0.3,
+        use_attention: bool = True,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.bidirectional = bidirectional
+        self.use_attention = use_attention
+        self.num_directions = 2 if bidirectional else 1
+
+        # Normalise input features
+        self.input_norm = nn.LayerNorm(input_dim)
+
+        # LSTM
+        self.lstm = nn.LSTM(
+            input_size=input_dim,
+            hidden_size=hidden_size,
+            num_layers=num_layers,
+            batch_first=True,
+            bidirectional=bidirectional,
+            dropout=dropout if num_layers > 1 else 0.0,
+        )
+
+        lstm_output_dim = hidden_size * self.num_directions
+
+        # Attention
+        if self.use_attention:
+            self.attention = SelfAttention(lstm_output_dim)
+
+        # Projection to output_dim
+        self.projection = nn.Sequential(
+            nn.Linear(lstm_output_dim, output_dim),
+            nn.BatchNorm1d(output_dim),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=dropout),
+        )
+
+    @classmethod
+    def from_config(cls, config) -> "SequenceAnalyzer":
+        """Construct from a Config object."""
+        sa = config.model.sequence_analyzer
+        fe = config.model.feature_extractor
+        return cls(
+            input_dim=fe.embedding_dim,
+            hidden_size=sa.hidden_size,
+            num_layers=sa.num_layers,
+            output_dim=fe.embedding_dim,
+            bidirectional=sa.bidirectional,
+            dropout=sa.dropout,
+            use_attention=sa.attention,
+        )
+
+    def forward(
+        self, features: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """
+        Args:
+            features: (B, T, D)  — sequence of CNN embeddings.
+
+        Returns:
+            encoding:          (B, output_dim) — temporal pattern encoding.
+            attention_weights: (B, T) or None  — per-timestep importance.
+        """
+        # Normalise
+        normed = self.input_norm(features)
+
+        # LSTM
+        lstm_out, _ = self.lstm(normed)  # (B, T, H*num_directions)
+
+        # Aggregate
+        if self.use_attention:
+            context, attn_weights = self.attention(lstm_out)
+        else:
+            # Fallback: use the last hidden state
+            context = lstm_out[:, -1, :]
+            attn_weights = None
+
+        # Project
+        encoding = self.projection(context)
+        return encoding, attn_weights