[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

neural_network/real_time_encoder_transformer.py

@@ -1,11 +1,11 @@
-
 from __future__ import annotations
 import math
 from typing import Optional, Tuple
 
 import numpy as np
 import pandas as pd
 
+
 def _softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
     x_max = np.max(x, axis=axis, keepdims=True)
     e = np.exp(x - x_max)
@@ -18,6 +18,7 @@ def _stable_div(x: np.ndarray, denom: np.ndarray) -> np.ndarray:
 
 # Time2Vec
 
+
 class Time2Vec:
     """
     Time2Vec positional encoding (simple) for real-valued time steps.
@@ -51,8 +52,15 @@ def forward(self, time_steps: np.ndarray) -> np.ndarray:
 
 # PositionwiseFeedForward
 
+
 class PositionwiseFeedForward:
-    def __init__(self, d_model: int, hidden: int, drop_prob: float = 0.0, seed: Optional[int] = None):
+    def __init__(
+        self,
+        d_model: int,
+        hidden: int,
+        drop_prob: float = 0.0,
+        seed: Optional[int] = None,
+    ):
         if seed is not None:
             np.random.seed(seed)
         # simple linear layers (no dropout during forward-only inference, but kept shape)
@@ -70,11 +78,17 @@ def forward(self, x: np.ndarray) -> np.ndarray:
         return out
 
 
-
 # Scaled Dot-Product Attention
 
+
 class ScaledDotProductAttention:
-    def forward(self, q: np.ndarray, k: np.ndarray, v: np.ndarray, mask: Optional[np.ndarray] = None) -> Tuple[np.ndarray, np.ndarray]:
+    def forward(
+        self,
+        q: np.ndarray,
+        k: np.ndarray,
+        v: np.ndarray,
+        mask: Optional[np.ndarray] = None,
+    ) -> Tuple[np.ndarray, np.ndarray]:
         """
         q,k,v: shapes (b, n_head, seq_len, d_k)
         mask: optional boolean or 0/1 mask of shape (b, seq_len) or (b, 1, 1, seq_len)
@@ -90,7 +104,11 @@ def forward(self, q: np.ndarray, k: np.ndarray, v: np.ndarray, mask: Optional[np
                 mask2 = mask[:, None, None, :]  # (b,1,1,seq_len)
             elif mask.ndim == 3:
                 # if provided as (b, n_head, seq_len) or (b, 1, seq_len)
-                mask2 = mask[:, None, :, :] if mask.shape[1] != seq_len else mask[:, None, None, :]
+                mask2 = (
+                    mask[:, None, :, :]
+                    if mask.shape[1] != seq_len
+                    else mask[:, None, None, :]
+                )
             else:
                 mask2 = mask
             # mask2==0 => masked
@@ -103,6 +121,7 @@ def forward(self, q: np.ndarray, k: np.ndarray, v: np.ndarray, mask: Optional[np
 
 # MultiHeadAttention
 
+
 class MultiHeadAttention:
     def __init__(self, d_model: int, n_head: int, seed: Optional[int] = None):
         if d_model % n_head != 0:
@@ -114,13 +133,21 @@ def __init__(self, d_model: int, n_head: int, seed: Optional[int] = None):
         self.d_k = d_model // n_head
 
         # weight matrices for q,k,v and output
-        self.w_q = np.random.randn(d_model, d_model) * math.sqrt(2.0 / (d_model + d_model))
+        self.w_q = np.random.randn(d_model, d_model) * math.sqrt(
+            2.0 / (d_model + d_model)
+        )
         self.b_q = np.zeros((d_model,))
-        self.w_k = np.random.randn(d_model, d_model) * math.sqrt(2.0 / (d_model + d_model))
+        self.w_k = np.random.randn(d_model, d_model) * math.sqrt(
+            2.0 / (d_model + d_model)
+        )
         self.b_k = np.zeros((d_model,))
-        self.w_v = np.random.randn(d_model, d_model) * math.sqrt(2.0 / (d_model + d_model))
+        self.w_v = np.random.randn(d_model, d_model) * math.sqrt(
+            2.0 / (d_model + d_model)
+        )
         self.b_v = np.zeros((d_model,))
-        self.w_out = np.random.randn(d_model, d_model) * math.sqrt(2.0 / (d_model + d_model))
+        self.w_out = np.random.randn(d_model, d_model) * math.sqrt(
+            2.0 / (d_model + d_model)
+        )
         self.b_out = np.zeros((d_model,))
 
         self.attn = ScaledDotProductAttention()
@@ -139,7 +166,13 @@ def _concat_heads(self, x: np.ndarray) -> np.ndarray:
         b, n_head, seq_len, d_k = x.shape
         return x.transpose(0, 2, 1, 3).reshape(b, seq_len, n_head * d_k)
 
-    def forward(self, query: np.ndarray, key: np.ndarray, value: np.ndarray, mask: Optional[np.ndarray] = None) -> Tuple[np.ndarray, np.ndarray]:
+    def forward(
+        self,
+        query: np.ndarray,
+        key: np.ndarray,
+        value: np.ndarray,
+        mask: Optional[np.ndarray] = None,
+    ) -> Tuple[np.ndarray, np.ndarray]:
         """
         query/key/value: (b, seq_len, d_model)
         returns: out (b, seq_len, d_model), attn_weights (b, n_head, seq_len, seq_len)
@@ -157,9 +190,9 @@ def forward(self, query: np.ndarray, key: np.ndarray, value: np.ndarray, mask: O
         return out, attn
 
 
-
 # LayerNorm
 
+
 class LayerNorm:
     def __init__(self, d_model: int, eps: float = 1e-12):
         self.gamma = np.ones((d_model,))
@@ -173,10 +206,14 @@ def forward(self, x: np.ndarray) -> np.ndarray:
         x_norm = (x - mean) / np.sqrt(var + self.eps)
         return self.gamma * x_norm + self.beta
 
+
 # TransformerEncoderLayer
 
+
 class TransformerEncoderLayer:
-    def __init__(self, d_model: int, n_head: int, hidden_dim: int, seed: Optional[int] = None):
+    def __init__(
+        self, d_model: int, n_head: int, hidden_dim: int, seed: Optional[int] = None
+    ):
         self.self_attn = MultiHeadAttention(d_model, n_head, seed=seed)
         self.ffn = PositionwiseFeedForward(d_model, hidden_dim, seed=seed)
         self.norm1 = LayerNorm(d_model)
@@ -193,26 +230,41 @@ def forward(self, x: np.ndarray, mask: Optional[np.ndarray] = None) -> np.ndarra
 
 # TransformerEncoder (stack)
 
+
 class TransformerEncoder:
-    def __init__(self, d_model: int, n_head: int, hidden_dim: int, num_layers: int, seed: Optional[int] = None):
-        self.layers = [TransformerEncoderLayer(d_model, n_head, hidden_dim, seed=seed) for _ in range(num_layers)]
+    def __init__(
+        self,
+        d_model: int,
+        n_head: int,
+        hidden_dim: int,
+        num_layers: int,
+        seed: Optional[int] = None,
+    ):
+        self.layers = [
+            TransformerEncoderLayer(d_model, n_head, hidden_dim, seed=seed)
+            for _ in range(num_layers)
+        ]
 
     def forward(self, x: np.ndarray, mask: Optional[np.ndarray] = None) -> np.ndarray:
         out = x
         for layer in self.layers:
             out = layer.forward(out, mask)
         return out
 
+
 # AttentionPooling
 
+
 class AttentionPooling:
     def __init__(self, d_model: int, seed: Optional[int] = None):
         if seed is not None:
             np.random.seed(seed)
         self.w = np.random.randn(d_model) * math.sqrt(2.0 / d_model)
         self.b = 0.0
 
-    def forward(self, x: np.ndarray, mask: Optional[np.ndarray] = None) -> Tuple[np.ndarray, np.ndarray]:
+    def forward(
+        self, x: np.ndarray, mask: Optional[np.ndarray] = None
+    ) -> Tuple[np.ndarray, np.ndarray]:
         """
         x: (b, seq_len, d_model)
         mask: (b, seq_len) where 1 = valid, 0 = pad
@@ -228,8 +280,10 @@ def forward(self, x: np.ndarray, mask: Optional[np.ndarray] = None) -> Tuple[np.
         pooled = np.matmul(weights[:, None, :], x).squeeze(1)  # (b, d_model)
         return pooled, weights
 
+
 # EEGTransformer (forward-only)
 
+
 class EEGTransformer:
     def __init__(
         self,
@@ -248,21 +302,29 @@ def __init__(
         self.d_model = d_model
         self.task_type = task_type
         # input projection
-        self.w_in = np.random.randn(feature_dim, d_model) * math.sqrt(2.0 / (feature_dim + d_model))
+        self.w_in = np.random.randn(feature_dim, d_model) * math.sqrt(
+            2.0 / (feature_dim + d_model)
+        )
         self.b_in = np.zeros((d_model,))
         # time embedding
         self.time2vec = Time2Vec(d_model, seed=seed)
-        self.encoder = TransformerEncoder(d_model, n_head, hidden_dim, num_layers, seed=seed)
+        self.encoder = TransformerEncoder(
+            d_model, n_head, hidden_dim, num_layers, seed=seed
+        )
         self.pooling = AttentionPooling(d_model, seed=seed)
         # output
-        self.w_out = np.random.randn(d_model, output_dim) * math.sqrt(2.0 / (d_model + output_dim))
+        self.w_out = np.random.randn(d_model, output_dim) * math.sqrt(
+            2.0 / (d_model + output_dim)
+        )
         self.b_out = np.zeros((output_dim,))
 
     def _input_proj(self, x: np.ndarray) -> np.ndarray:
         # x: (b, seq_len, feature_dim) -> (b, seq_len, d_model)
         return np.tensordot(x, self.w_in, axes=([2], [0])) + self.b_in
 
-    def forward(self, x: np.ndarray, mask: Optional[np.ndarray] = None) -> Tuple[np.ndarray, np.ndarray]:
+    def forward(
+        self, x: np.ndarray, mask: Optional[np.ndarray] = None
+    ) -> Tuple[np.ndarray, np.ndarray]:
         """
         x: (b, seq_len, feature_dim)
         mask: optional (b, seq_len) 1=valid,0=pad
@@ -276,7 +338,9 @@ def forward(self, x: np.ndarray, mask: Optional[np.ndarray] = None) -> Tuple[np.
         x_proj = self._input_proj(x) + time_emb  # broadcast add -> (b,t,d_model)
         enc = self.encoder.forward(x_proj, mask)
         pooled, attn_weights = self.pooling.forward(enc, mask)
-        out = np.tensordot(pooled, self.w_out, axes=([1], [0])) + self.b_out  # (b,output_dim)
+        out = (
+            np.tensordot(pooled, self.w_out, axes=([1], [0])) + self.b_out
+        )  # (b,output_dim)
         if self.task_type == "classification":
             out = _softmax(out, axis=-1)
         return out, attn_weights
@@ -292,7 +356,15 @@ def forward(self, x: np.ndarray, mask: Optional[np.ndarray] = None) -> Tuple[np.
     rng = np.random.RandomState(42)
     X = rng.randn(batch, seq_len, feature_dim).astype(float)
 
-    model = EEGTransformer(feature_dim=feature_dim, d_model=32, n_head=4, hidden_dim=64, num_layers=2, output_dim=1, seed=0)
+    model = EEGTransformer(
+        feature_dim=feature_dim,
+        d_model=32,
+        n_head=4,
+        hidden_dim=64,
+        num_layers=2,
+        output_dim=1,
+        seed=0,
+    )
     out, attn_weights = model.forward(X)
     print("Output shape:", out.shape)
     print("Output:", out)