update flake8 and format

Author: loki-veera

.flake8

@@ -0,0 +1,45 @@
+#########################
+# Flake8 Configuration  #
+# (.flake8)             #
+#########################
+[flake8]
+ignore =
+    # asserts are ok when testing.
+    S101
+    # pickle
+    S301
+    # pickle
+    S403
+    S404
+    S603
+    # Line break before binary operator (flake8 is wrong)
+    W503
+    # Ignore the spaces black puts before columns.
+    E203
+    # allow path extensions for testing.
+    E402
+    DAR101
+    DAR201
+    # flake and pylance disagree on linebreaks in strings.
+    N400
+    N806
+exclude =
+    .tox,
+    .git,
+    __pycache__,
+    docs/source/conf.py,
+    build,
+    dist,
+    tests/fixtures/*,
+    *.pyc,
+    *.bib,
+    *.egg-info,
+    .cache,
+    .eggs,
+    data.
+max-line-length = 120
+max-complexity = 20
+import-order-style = pycharm
+application-import-names =
+    seleqt
+    tests

.github/workflows/tests.yml

@@ -0,0 +1,33 @@
+name: Tests
+on: [ push, pull_request ]
+jobs:
+  tests:
+    name: Tests
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install uv
+        uses: astral-sh/setup-uv@v3
+      - name: Set up Python 3.8.20
+        run: uv python install 3.8.20
+      - name: Install dependencies (CPU only)
+        run: uv sync --extra cpu --no-dev
+      - name: Test with pytest
+        run: uv run pytest
+
+  lint:
+    name: Lint
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.8.20"]
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install uv
+        uses: astral-sh/setup-uv@v3
+      - name: Set up Python ${{ matrix.python-version }}
+        run: uv python install ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: uv tool install flake8
+      - name: Run flake8
+        run: uv tool run flake8 src/ tests/

.pre-commit-config.yaml

@@ -0,0 +1,33 @@
+repos:
+  - repo: local
+    hooks:
+      - id: isort
+        name: isort
+        entry: uv run isort
+        language: system
+        args: ["--profile", "black"]
+        files: ^(src/|tests/)
+        types: [python]
+
+      - id: black  
+        name: black
+        entry: uv run black
+        language: system 
+        args: ["--target-version", "py312"]
+        files: ^(src/|tests/)
+        types: [python]
+
+      - id: flake8
+        name: flake8
+        entry: uv run flake8
+        language: system
+        files: ^(src/|tests/)$
+        types: [python]
+
+      - id: pytest
+        name: pytest
+        entry: uv run pytest
+        language: system
+        files: ^(tests/)$
+        pass_filenames: false
+        always_run: true

Makefile

@@ -0,0 +1,24 @@
+.PHONY: format check-format test lint check-all typing all
+
+# Format code
+format:
+	uv run isort tests/ src/
+	uv run black tests/ src/
+
+# Check formatting (without making changes)  
+check-format:
+	uv run isort --check-only --diff tests/ src/
+	uv run black --check --diff tests/ src/
+
+lint:
+	uv run flake8 tests/ src/
+
+test:
+	uv run pytest
+
+# Run all pre-commit checks manually
+check-all:
+	uv run pre-commit run --all-files
+
+all: format test lint
+	@echo "All checks passed!"

src/cara/__init__.py

@@ -1 +1 @@
-"""src package init."""
+"""src package init."""

src/cara/cara.py

@@ -1,25 +1,42 @@
-from copy import deepcopy
-from typing import Dict, Any
+"""Implement CaRA."""
 
+from typing import Any, Dict
+
+import tensorly as tl
+import timm
 import torch as th
 import torch.nn as nn
-import timm
-import tensorly as tl
+
 tl.set_backend("pytorch")
 
+global_model: th.nn.Module
 
-def cp_attn(self, x):
+
+def cp_attn(self, x: th.Tensor) -> th.Tensor:
+    """Attention with CP parameters.
+
+    Args:
+        x (th.Tensor): Input tensor.
+
+    Returns:
+        th.Tensor: CaRA attention output.
+    """
     B, N, C = x.shape
     qkv = self.qkv(x)
-    f1 = global_model.CP_A1[self.attn_idx:self.attn_idx+3, :]
-    tensor_attn = tl.cp_to_tensor((global_model.CP_R1, (f1, global_model.CP_A2, global_model.CP_A3, global_model.CP_A4)))
+    f1 = global_model.CP_A1[self.attn_idx : self.attn_idx + 3, :]
+    tensor_attn = tl.cp_to_tensor(
+        (
+            global_model.CP_R1,
+            (f1, global_model.CP_A2, global_model.CP_A3, global_model.CP_A4),
+        )
+    )
     K, E, H, D = tensor_attn.shape
-    tensor_attn = tensor_attn.reshape((K, E, H*D))
+    tensor_attn = tensor_attn.reshape((K, E, H * D))
     qkv_delta = th.einsum("bnd, kde->kbne", x, self.dp(tensor_attn))
-    qkv_delta = qkv_delta.reshape(3, B, N, self.num_heads, C//self.num_heads).permute(
-        0, 1, 3, 2, 4
-    )
-    qkv = qkv.reshape(B, N, 3, self.num_heads, C//self.num_heads).permute(
+    qkv_delta = qkv_delta.reshape(
+        3, B, N, self.num_heads, C // self.num_heads
+    ).permute(0, 1, 3, 2, 4)
+    qkv = qkv.reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(
         2, 0, 3, 1, 4
     )
     qkv += qkv_delta * self.s
@@ -28,56 +45,83 @@ def cp_attn(self, x):
     attn = attn.softmax(dim=-1)
     attn = self.attn_drop(attn)
 
-    x = (attn@v).transpose(1, 2).reshape(B, N, C)
+    x = (attn @ v).transpose(1, 2).reshape(B, N, C)
 
     proj = self.proj(x)
-    p1 = global_model.CP_P1[self.idx:self.idx+1, :]
-    tensor_proj = tl.cp_to_tensor((global_model.CP_R2, (p1, global_model.CP_P2, global_model.CP_P3)))
+    p1 = global_model.CP_P1[self.idx : self.idx + 1, :]
+    tensor_proj = tl.cp_to_tensor(
+        (global_model.CP_R2, (p1, global_model.CP_P2, global_model.CP_P3))
+    )
     AA, AB, AC = tensor_proj.shape
-    tensor_proj = tensor_proj.reshape((AA*AB, AC))
-    proj_delta = x@self.dp(tensor_proj.T) + global_model.CP_bias1
+    tensor_proj = tensor_proj.reshape((AA * AB, AC))
+    proj_delta = x @ self.dp(tensor_proj.T) + global_model.CP_bias1
     proj += proj_delta * self.s
     x = self.proj_drop(proj)
     return x
 
 
-def cp_mlp(self, x):
-    p1_up = global_model.CP_P1[self.idx:self.idx+4, :]
-    p1_down = global_model.CP_P1[self.idx+4: self.idx+8, :]
+def cp_mlp(self, x: th.Tensor) -> th.Tensor:
+    """Mlp with CP parameters.
+
+    Args:
+        x (th.Tensor): Input tensor.
+
+    Returns:
+        th.Tensor: Mlp projected output.
+    """
+    p1_up = global_model.CP_P1[self.idx : self.idx + 4, :]
+    p1_down = global_model.CP_P1[self.idx + 4 : self.idx + 8, :]
 
     up = self.fc1(x)
-    tensor_up = tl.cp_to_tensor((global_model.CP_R2, (p1_up, global_model.CP_P2, global_model.CP_P3)))
+    tensor_up = tl.cp_to_tensor(
+        (global_model.CP_R2, (p1_up, global_model.CP_P2, global_model.CP_P3))
+    )
     AA, AB, AC = tensor_up.shape
-    tensor_up = tensor_up.reshape((AA*AB, AC))
-    up_delta = x@self.dp(tensor_up.T) + global_model.CP_bias2
+    tensor_up = tensor_up.reshape((AA * AB, AC))
+    up_delta = x @ self.dp(tensor_up.T) + global_model.CP_bias2
     up += up_delta * self.s
 
     x = self.act(up)
     x = self.drop(x)
-    
+
     down = self.fc2(x)
-    tensor_down = tl.cp_to_tensor((global_model.CP_R2, (p1_down, global_model.CP_P2, global_model.CP_P3)))
-    tensor_down = tensor_down.reshape((AA*AB, AC))
-    down_delta = x@self.dp(tensor_down) + global_model.CP_bias3
+    tensor_down = tl.cp_to_tensor(
+        (global_model.CP_R2, (p1_down, global_model.CP_P2, global_model.CP_P3))
+    )
+    tensor_down = tensor_down.reshape((AA * AB, AC))
+    down_delta = x @ self.dp(tensor_down) + global_model.CP_bias3
     down += down_delta * self.s
     x = self.drop(down)
     return x
 
 
-def set_cara(model: nn.Module, rank: int, scale: float, l_mu: float, l_std: float):
-    if type(model) == timm.models.vision_transformer.VisionTransformer:
+def set_cara(
+    model: nn.Module, rank: int, scale: float, l_mu: float, l_std: float
+) -> None:
+    """Cara setup.
+
+    Args:
+        model (nn.Module): ViT model.
+        rank (int): FT Rank.
+        scale (float): FT scale.
+        l_mu (float): Init lambda_mu.
+        l_std (float): Init lambda_std.
+    """
+    if type(model) is timm.models.vision_transformer.VisionTransformer:
         # Declare CaRA parameters
         model.CP_A1 = nn.Parameter(th.empty([36, rank]), requires_grad=True)
         model.CP_A2 = nn.Parameter(th.empty([768, rank]), requires_grad=True)
         model.CP_A3 = nn.Parameter(th.empty([12, rank]), requires_grad=True)
-        model.CP_A4 = nn.Parameter(th.empty([768//12, rank]), requires_grad=True)
+        model.CP_A4 = nn.Parameter(
+            th.empty([768 // 12, rank]), requires_grad=True
+        )
         model.CP_P1 = nn.Parameter(th.empty([108, rank]), requires_grad=True)
         model.CP_P2 = nn.Parameter(th.empty([768, rank]), requires_grad=True)
         model.CP_P3 = nn.Parameter(th.empty([768, rank]), requires_grad=True)
         model.CP_R1 = nn.Parameter(th.empty([rank]), requires_grad=True)
         model.CP_R2 = nn.Parameter(th.empty([rank]), requires_grad=True)
         model.CP_bias1 = nn.Parameter(th.empty([768]), requires_grad=True)
-        model.CP_bias2 = nn.Parameter(th.empty([768*4]), requires_grad=True)
+        model.CP_bias2 = nn.Parameter(th.empty([768 * 4]), requires_grad=True)
         model.CP_bias3 = nn.Parameter(th.empty([768]), requires_grad=True)
         # Initialise CaRA parameters
         nn.init.xavier_normal_(model.CP_A1)
@@ -100,7 +144,7 @@ def set_cara(model: nn.Module, rank: int, scale: float, l_mu: float, l_std: floa
         model.idx = 0
         model.attn_idx = 0
     for child in model.children():
-        if type(child) == timm.models.vision_transformer.Attention:
+        if type(child) is timm.models.vision_transformer.Attention:
             child.dp = nn.Dropout(0.1)
             child.s = scale
             child.dim = rank
@@ -109,28 +153,36 @@ def set_cara(model: nn.Module, rank: int, scale: float, l_mu: float, l_std: floa
             global_model.idx += 1
             global_model.attn_idx += 3
             bound_method = cp_attn.__get__(child, child.__class__)
-            setattr(child, "forward", bound_method)
-        elif type(child) == timm.models.layers.mlp.Mlp:
+            setattr(child, "forward", bound_method)  # noqa: B010
+        elif type(child) is timm.models.layers.mlp.Mlp:
             child.dp = nn.Dropout(0.1)
             child.s = scale
             child.dim = rank
             child.idx = global_model.idx
             global_model.idx += 8
             bound_method = cp_mlp.__get__(child, child.__class__)
-            setattr(child, "forward", bound_method)
+            setattr(child, "forward", bound_method)  # noqa: B010
         elif len(list(child.children())) != 0:
             set_cara(child, rank, scale, l_mu, l_std)
-            
 
-def cara(config):
+
+def cara(config: Dict[str, Any]) -> th.nn.Module:
+    """Set CaRA for the given configuration.
+
+    Args:
+        config (Dict[str, Any]): Dictionary containing CaRA configuration.
+
+    Returns:
+        th.nn.Module: CaRA model.
+    """
     # CaRA parameters
     model = config["model"]
     rank = config["rank"]
     scale = config["scale"]
     l_mu = config["l_mu"]
     l_std = config["l_std"]
-    
+
     global global_model
     global_model = model
     set_cara(model, rank, scale, l_mu, l_std)
-    return global_model
+    return global_model

tests/__init__.py

@@ -1 +1 @@
-"""Test package."""
+"""Test package."""