sawnergy 1.0.3__py3-none-any.whl → 1.0.9__py3-none-any.whl

sawnergy/embedding/SGNS_torch.py

@@ -27,46 +27,77 @@ class SGNS_Torch:
     def __init__(self,
                  V: int,
                  D: int,
+                 in_weights: torch.Tensor | np.ndarray | None = None,
+                 out_weights: torch.Tensor | np.ndarray | None = None,
                  *,
-                seed: int | None = None,
-                optim: Type[Optimizer],
-                optim_kwargs: dict,
-                lr_sched: Type[LRScheduler] | None = None,
-                lr_sched_kwargs: dict | None = None,
-                device: str | None = None):
-        """
+                 seed: int | None = None,
+                 optim: Type[Optimizer] = torch.optim.SGD,
+                 optim_kwargs: dict | None = None,
+                 lr_sched: Type[LRScheduler] | None = None,
+                 lr_sched_kwargs: dict | None = None,
+                 device: str | None = None):
+        """Initialize SGNS (negative sampling) in PyTorch.
+
+        Shapes:
+            - Embedding tables:
+                in_weights:  (V, D) or None — row i is the “input” vector for token i.
+                out_weights: (V, D) or None — row i is the “output” vector for token i.
+
         Args:
-            V: Vocabulary size (number of nodes).
+            V: Vocabulary size (number of nodes/tokens).
             D: Embedding dimensionality.
-            seed: Optional RNG seed for PyTorch.
-            optim: Optimizer class to instantiate.
-            optim_kwargs: Keyword arguments for the optimizer.
+            in_weights: Optional starting input-embedding matrix of shape (V, D).
+            out_weights: Optional starting output-embedding matrix of shape (V, D).
+            seed: Optional RNG seed for PyTorch (controls init, sampling, and shuffles).
+            optim: Optimizer class to instantiate. Defaults to plain SGD.
+            optim_kwargs: Keyword arguments for the optimizer. Defaults to {"lr": 0.1}.
             lr_sched: Optional learning-rate scheduler class.
-            lr_sched_kwargs: Keyword arguments for the scheduler.
-            device: Target device string (e.g. ``"cuda"``). Defaults to CUDA if available, else CPU.
+            lr_sched_kwargs: Keyword arguments for the scheduler (required if lr_sched is provided).
+            device: Target device string (e.g. "cuda"). Defaults to CUDA if available, else CPU.
         """
-        if optim_kwargs is None:
-            raise ValueError("optim_kwargs must be provided")
+        optim_kwargs = optim_kwargs or {"lr": 0.1}
         if lr_sched is not None and lr_sched_kwargs is None:
             raise ValueError("lr_sched_kwargs required when lr_sched is provided")
+
         self.V, self.D = int(V), int(D)
-        resolved_device = device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu")
+
+        resolved_device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         self.device = torch.device(resolved_device)
-        _logger.info("SGNS_Torch init: V=%d D=%d device=%s seed=%s", self.V, self.D, self.device, seed)
 
-        if seed is not None:
-            torch.manual_seed(int(seed))
-            np.random.seed(int(seed))
+        # Seed torch
+        self.seed = None if seed is None else int(seed)
+        if self.seed is not None:
+            torch.manual_seed(self.seed)
             if self.device.type == "cuda":
-                torch.cuda.manual_seed_all(int(seed))
+                torch.cuda.manual_seed_all(self.seed)
 
         # two embeddings as in/out matrices
-        self.in_emb  = nn.Embedding(self.V, self.D)
-        self.out_emb = nn.Embedding(self.V, self.D)
+        self.in_emb  = nn.Embedding(self.V, self.D, device=self.device)
+        self.out_emb = nn.Embedding(self.V, self.D, device=self.device)
+
+        # init / warm-start
+        with torch.no_grad():
+            if in_weights is not None:
+                w = torch.as_tensor(in_weights, dtype=torch.float32, device=self.device)
+                if w.shape != (self.V, self.D):
+                    raise ValueError(f"in_weights must be (V,D); got {tuple(w.shape)}")
+                self.in_emb.weight.copy_(w)
+            else:
+                nn.init.uniform_(self.in_emb.weight, -0.5 / self.D, 0.5 / self.D)
+
+            if out_weights is not None:
+                w = torch.as_tensor(out_weights, dtype=torch.float32, device=self.device)
+                if w.shape != (self.V, self.D):
+                    raise ValueError(f"out_weights must be (V,D); got {tuple(w.shape)}")
+                self.out_emb.weight.copy_(w)
+            else:
+                nn.init.zeros_(self.out_emb.weight)
 
         self.to(self.device)
+        _logger.info("SGNS_Torch init: V=%d D=%d device=%s seed=%s", self.V, self.D, self.device, self.seed)
 
         params = list(self.in_emb.parameters()) + list(self.out_emb.parameters())
+        # optimizer / scheduler
         self.opt = optim(params=params, **optim_kwargs)
         self.lr_sched = lr_sched(self.opt, **lr_sched_kwargs) if lr_sched is not None else None
 
@@ -74,7 +105,17 @@ class SGNS_Torch:
                 center: torch.Tensor,
                 pos: torch.Tensor,
                 neg: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute positive/negative logits for SGNS.
 
+        Inputs:
+            center: int tensor of shape (B,), values in [0, V)
+            pos:    int tensor of shape (B,), values in [0, V)
+            neg:    int tensor of shape (B, K), values in [0, V)
+
+        Returns:
+            pos_logits: (B,)
+            neg_logits: (B, K)
+        """
         center = center.to(self.device, dtype=torch.long)
         pos    = pos.to(self.device, dtype=torch.long)
         neg    = neg.to(self.device, dtype=torch.long)
@@ -83,9 +124,8 @@ class SGNS_Torch:
         pe = self.out_emb(pos)    # (B, D)
         ne = self.out_emb(neg)    # (B, K, D)
 
-        pos_logits = (c * pe).sum(dim=-1)              # (B,)
-        neg_logits = (c.unsqueeze(1) * ne).sum(dim=-1) # (B, K)
-
+        pos_logits = (c * pe).sum(dim=-1)               # (B,)
+        neg_logits = (c.unsqueeze(1) * ne).sum(dim=-1)  # (B, K)
         return pos_logits, neg_logits
 
     __call__ = predict
@@ -112,15 +152,27 @@ class SGNS_Torch:
         idx = np.arange(N)
 
         noise_probs = torch.as_tensor(noise_dist, dtype=torch.float32, device=self.device)
+        # normalize if slightly off; enforce nonnegativity + finite sum
+        if (noise_probs < 0).any():
+            raise ValueError("noise_dist has negative entries")
+        s = noise_probs.sum()
+        if not torch.isfinite(s) or float(s.item()) <= 0.0:
+            raise ValueError("noise_dist must have positive finite sum")
+        if abs(float(s.item()) - 1.0) > 1e-6:
+            noise_probs = noise_probs / s
 
         for epoch in range(1, int(num_epochs) + 1):
             epoch_loss = 0.0
             batches = 0
+
             if shuffle_data:
-                np.random.shuffle(idx)
+                if self.seed is None:
+                    np.random.shuffle(idx)
+                else:
+                    np.random.default_rng(self.seed + epoch).shuffle(idx)
 
-            for s in range(0, N, int(batch_size)):
-                take = idx[s:s+int(batch_size)]
+            for s_ in range(0, N, int(batch_size)):
+                take = idx[s_:s_+int(batch_size)]
                 if take.size == 0:
                     continue
                 K = int(num_negative_samples)
@@ -128,19 +180,199 @@ class SGNS_Torch:
 
                 cen = torch.as_tensor(centers[take],  dtype=torch.long, device=self.device)  # (B,)
                 pos = torch.as_tensor(contexts[take], dtype=torch.long, device=self.device)  # (B,)
-                neg = torch.multinomial(noise_probs, num_samples=B * K, replacement=True).view(B, K)  # (B,K) on device
+                neg = torch.multinomial(noise_probs, num_samples=B * K, replacement=True).view(B, K)  # (B,K)
 
                 pos_logits, neg_logits = self(cen, pos, neg)
 
-                # BCE(+)
                 y_pos = torch.ones_like(pos_logits)
-                loss_pos = bce(pos_logits, y_pos)
-
-                # BCE(-):
                 y_neg = torch.zeros_like(neg_logits)
+                loss_pos = bce(pos_logits, y_pos)
                 loss_neg = bce(neg_logits, y_neg)
 
-                loss = loss_pos + loss_neg
+                loss = loss_pos + K * loss_neg
+
+                self.opt.zero_grad(set_to_none=True)
+                loss.backward()
+                self.opt.step()
+
+                if lr_step_per_batch and self.lr_sched is not None:
+                    self.lr_sched.step()
+
+                epoch_loss += float(loss.detach().cpu().item())
+                batches += 1
+                _logger.debug("Epoch %d batch %d loss=%.6f", epoch, batches, loss.item())
+
+            if not lr_step_per_batch and self.lr_sched is not None:
+                self.lr_sched.step()
+
+            mean_loss = epoch_loss / max(batches, 1)
+            _logger.info("Epoch %d/%d mean_loss=%.6f", epoch, num_epochs, mean_loss)
+
+    @property
+    def in_embeddings(self) -> np.ndarray:
+        W = self.in_emb.weight.detach().cpu().numpy()  # (V, D)
+        _logger.debug("In emb shape: %s", W.shape)
+        return W
+
+    @property
+    def out_embeddings(self) -> np.ndarray:
+        W = self.out_emb.weight.detach().cpu().numpy()  # (V, D)
+        _logger.debug("Out emb shape: %s", W.shape)
+        return W
+
+    @property
+    def avg_embeddings(self) -> np.ndarray:
+        return 0.5 * (self.in_embeddings + self.out_embeddings)
+
+    # tiny helper for device move
+    def to(self, device):
+        self.in_emb.to(device)
+        self.out_emb.to(device)
+        return self
+
+
+class SG_Torch:
+    """PyTorch implementation of Skip-Gram (full softmax, **no biases**).
+
+    This variant uses **no bias terms**: both projections are pure linear maps.
+
+    Computation:
+        x = one_hot(center, V)          # (B, V)
+        y = x @ W_in                    # (B, D), with W_in ∈ R^{VxD}
+        logits = y @ W_out              # (B, V), with W_out ∈ R^{DxV}
+        loss = CrossEntropyLoss(logits, context)
+
+    Embeddings:
+        - Input embeddings  = rows of W_in        → shape (V, D)
+        - Output embeddings = rows of W_outᵀ      → shape (V, D)
+    """
+
+    def __init__(self,
+                 V: int,
+                 D: int,
+                 in_weights: torch.Tensor | np.ndarray | None = None,
+                 out_weights: torch.Tensor | np.ndarray | None = None,
+                 *,
+                 seed: int | None = None,
+                 optim: Type[Optimizer] = torch.optim.SGD,
+                 optim_kwargs: dict | None = None,
+                 lr_sched: Type[LRScheduler] | None = None,
+                 lr_sched_kwargs: dict | None = None,
+                 device: str | None = None):
+        """Initialize the plain Skip-Gram (full softmax, **no biases**) model in PyTorch.
+
+        Shapes:
+            - Linear maps (no bias):
+                W_in:  (V, D) — rows are input embeddings for tokens.
+                W_out: (D, V) — maps D→V; rows of W_outᵀ are output embeddings.
+
+            - Warm-starts:
+                in_weights:  (V, D) or None — copied into W_in if provided.
+                out_weights: (D, V) or None — copied into W_out if provided.
+
+        Args:
+            V: Vocabulary size (number of nodes/tokens).
+            D: Embedding dimensionality.
+            in_weights: Optional starting matrix for W_in with shape (V, D).
+            out_weights: Optional starting matrix for W_out with shape (D, V).
+            seed: Optional RNG seed for reproducibility.
+            optim: Optimizer class to instantiate. Defaults to :class:`torch.optim.SGD`.
+            optim_kwargs: Keyword args for the optimizer. Defaults to ``{"lr": 0.1}``.
+            lr_sched: Optional learning-rate scheduler class.
+            lr_sched_kwargs: Keyword args for the scheduler (required if ``lr_sched`` is provided).
+            device: Target device string (e.g., ``"cuda"``). Defaults to CUDA if available, else CPU.
+
+        Notes:
+            The encoder/decoder are **bias-free** linear layers acting on one-hot centers:
+            • ``in_emb = nn.Linear(V, D, bias=False)``
+            • ``out_emb = nn.Linear(D, V, bias=False)``
+            Forward pass produces vocabulary-sized logits and is trained with CrossEntropyLoss.
+        """
+        optim_kwargs = optim_kwargs or {"lr": 0.1}
+        if lr_sched is not None and lr_sched_kwargs is None:
+            raise ValueError("lr_sched_kwargs required when lr_sched is provided")
+
+        self.V, self.D = int(V), int(D)
+
+        resolved_device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = torch.device(resolved_device)
+
+        # Seed torch (no global NumPy seeding)
+        self.seed = None if seed is None else int(seed)
+        if self.seed is not None:
+            torch.manual_seed(self.seed)
+            if self.device.type == "cuda":
+                torch.cuda.manual_seed_all(self.seed)
+
+        self.in_emb  = nn.Linear(self.V, self.D, bias=False, device=self.device)
+        self.out_emb = nn.Linear(self.D, self.V, bias=False, device=self.device)
+
+        # warm-starts (note Linear weights are (out_features, in_features))
+        with torch.no_grad():
+            if in_weights is not None:
+                w_in = torch.as_tensor(in_weights, dtype=torch.float32, device=self.device)
+                if w_in.shape != (self.V, self.D):
+                    raise ValueError(f"in_weights must be (V,D); got {tuple(w_in.shape)}")
+                self.in_emb.weight.copy_(w_in.T)  # (D,V)
+            # else: use default PyTorch init
+
+            if out_weights is not None:
+                w_out = torch.as_tensor(out_weights, dtype=torch.float32, device=self.device)
+                if w_out.shape != (self.D, self.V):
+                    raise ValueError(f"out_weights must be (D,V); got {tuple(w_out.shape)}")
+                self.out_emb.weight.copy_(w_out)  # (V,D) weight is (V,D) because (out=in V, in=D)
+            # else: default init
+
+        self.to(self.device)
+        _logger.info("SG_Torch init: V=%d D=%d device=%s seed=%s", self.V, self.D, self.device, self.seed)
+
+        params = list(self.in_emb.parameters()) + list(self.out_emb.parameters())
+        # optimizer / scheduler
+        self.opt = optim(params=params, **optim_kwargs)
+        self.lr_sched = lr_sched(self.opt, **lr_sched_kwargs) if lr_sched is not None else None
+
+    def predict(self, center: torch.Tensor) -> torch.Tensor:
+        center = center.to(self.device, dtype=torch.long)
+        c = nn.functional.one_hot(center, num_classes=self.V).to(dtype=torch.float32, device=self.device)
+        y = self.in_emb(c)
+        z = self.out_emb(y)
+        return z
+
+    __call__ = predict
+
+    def fit(self,
+            centers: np.ndarray,
+            contexts: np.ndarray,
+            num_epochs: int,
+            batch_size: int,
+            shuffle_data: bool,
+            lr_step_per_batch: bool,
+            **_ignore):
+        cce = nn.CrossEntropyLoss(reduction="mean")
+
+        N = centers.shape[0]
+        idx = np.arange(N)
+
+        for epoch in range(1, int(num_epochs) + 1):
+            epoch_loss = 0.0
+            batches = 0
+
+            if shuffle_data:
+                if self.seed is None:
+                    np.random.shuffle(idx)
+                else:
+                    np.random.default_rng(self.seed + epoch).shuffle(idx)
+
+            for s in range(0, N, int(batch_size)):
+                take = idx[s:s+int(batch_size)]
+                if take.size == 0:
+                    continue
+
+                cen = torch.as_tensor(centers[take], dtype=torch.long, device=self.device)
+                ctx = torch.as_tensor(contexts[take], dtype=torch.long, device=self.device)
+
+                logits = self(cen)
+                loss = cce(logits, ctx)
 
                 self.opt.zero_grad(set_to_none=True)
                 loss.backward()
@@ -160,9 +392,20 @@ class SGNS_Torch:
             _logger.info("Epoch %d/%d mean_loss=%.6f", epoch, num_epochs, mean_loss)
 
     @property
-    def embeddings(self) -> np.ndarray:
-        """Return the input embedding matrix as a NumPy array."""
-        return self.in_emb.weight.detach().cpu().numpy()
+    def in_embeddings(self) -> np.ndarray:
+        W = self.in_emb.weight.detach().T.cpu().numpy()  # (V, D)
+        _logger.debug("In emb shape: %s", W.shape)
+        return W
+
+    @property
+    def out_embeddings(self) -> np.ndarray:
+        W = self.out_emb.weight.detach().cpu().numpy()   # (V, D)
+        _logger.debug("Out emb shape: %s", W.shape)
+        return W
+
+    @property
+    def avg_embeddings(self) -> np.ndarray:
+        return 0.5 * (self.in_embeddings + self.out_embeddings)
 
     # tiny helper for device move
     def to(self, device):
@@ -171,7 +414,7 @@ class SGNS_Torch:
         return self
 
 
-__all__ = ["SGNS_Torch"]
+__all__ = ["SGNS_Torch", "SG_Torch"]
 
 if __name__ == "__main__":
     pass

sawnergy/embedding/__init__.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from .embedder import Embedder
+from .embedder import Embedder, align_frames
+from .visualizer import Visualizer
 
 def __getattr__(name: str):
     """Lazily expose optional backends."""
@@ -14,6 +15,16 @@ def __getattr__(name: str):
             ) from exc
         return SGNS_Torch
 
+    if name == "SG_Torch":
+        try:
+            from .SGNS_torch import SG_Torch
+        except Exception as exc:
+            raise ImportError(
+                "PyTorch backend requested but torch is not installed. "
+                "Install PyTorch via `pip install torch` (see https://pytorch.org/get-started)."
+            ) from exc
+        return SG_Torch
+
     if name == "SGNS_PureML":
             try:
                 from .SGNS_pml import SGNS_PureML
@@ -24,11 +35,25 @@ def __getattr__(name: str):
                     "Install PureML first via `pip install ym-pure-ml` "
                 ) from exc
 
+    if name == "SG_PureML":
+            try:
+                from .SGNS_pml import SG_PureML
+                return SG_PureML
+            except Exception as exc:
+                raise ImportError(
+                    "PureML is not installed. "
+                    "Install PureML first via `pip install ym-pure-ml` "
+                ) from exc
+
     raise AttributeError(name)
 
 
 __all__ = [
     "Embedder",
+    "align_frames",
+    "Visualizer",
     "SGNS_PureML",
     "SGNS_Torch",
+    "SG_PureML",
+    "SG_Torch"
 ]