cogcore/backbone/shared.py

"""English implementation note."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import numpy as np

@dataclass
class BackboneConfig:
    in_dim: int = 2048
    out_dim: int = 128
    seed: int = 7

class SharedBackbone:
    """English implementation note."""

    def __init__(self, cfg: BackboneConfig | None=None):
        self.cfg = cfg or BackboneConfig()
        rng = np.random.default_rng(self.cfg.seed)
        scale = float(np.sqrt(2.0 / max(self.cfg.out_dim, 1)))
        self.W = rng.normal(0.0, scale, size=(self.cfg.in_dim, self.cfg.out_dim)).astype(np.float32)

    def encode(self, X: np.ndarray) -> np.ndarray:
        """English implementation note."""
        return (X.astype(np.float32) @ self.W).astype(np.float32)

    def save(self, path: str | Path) -> None:
        np.savez_compressed(path, W=self.W, in_dim=self.cfg.in_dim, out_dim=self.cfg.out_dim, seed=self.cfg.seed)

    @classmethod
    def load(cls, path: str | Path) -> 'SharedBackbone':
        data = np.load(path)
        cfg = BackboneConfig(in_dim=int(data['in_dim']), out_dim=int(data['out_dim']), seed=int(data['seed']))
        obj = cls(cfg)
        obj.W = data['W'].astype(np.float32)
        return obj