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

sawnergy/__init__.py

@@ -3,11 +3,13 @@ from . import logging_util
 from . import rin
 from . import visual
 from . import walks
+from . import embedding
 
 __all__ = [
     "sawnergy_util",
     "logging_util",
     "rin",
     "visual",
-    "walks"
+    "walks",
+    "embedding"
 ]

sawnergy/embedding/SGNS_pml.py

@@ -3,11 +3,11 @@ from __future__ import annotations
 # third party
 import numpy as np
 from pureml.machinery import Tensor
-from pureml.layers import Embedding
-from pureml.losses import BCE
+from pureml.layers import Embedding, Affine
+from pureml.losses import BCE, CCE
 from pureml.general_math import sum as t_sum
-from pureml.optimizers import Optim, LRScheduler
-from pureml.training_utils import TensorDataset, DataLoader
+from pureml.optimizers import Optim, LRScheduler, SGD
+from pureml.training_utils import TensorDataset, DataLoader, one_hot
 from pureml.base import NN
 
 # built-in
@@ -28,51 +28,98 @@ class SGNS_PureML(NN):
     """PureML implementation of Skip-Gram with Negative Sampling."""
 
     def __init__(self,
-                V: int,
-                D: int,
-                *,
-                seed: int | None = None,
-                optim: Type[Optim],
-                optim_kwargs: dict,
-                lr_sched: Type[LRScheduler],
-                lr_sched_kwargs: dict):
+                 V: int,
+                 D: int,
+                 in_weights:  Tensor | np.ndarray | None = None,
+                 out_weights: Tensor | np.ndarray | None = None,
+                 *,
+                 seed: int | None = None,
+                 optim: Type[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.
+
+        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 negative sampling.
-            optim: PureML optimizer class.
-            optim_kwargs: Keyword arguments forwarded to the optimizer.
-            lr_sched: PureML learning-rate scheduler class.
-            lr_sched_kwargs: Keyword arguments forwarded to the scheduler.
+            in_weights: Optional starting input-embedding matrix of shape (V, D) as
+                        :class:`Tensor` or :class:`numpy.ndarray`. If None, the Embedding
+                        layer initializes it (seeded if `seed` is set).
+            out_weights: Optional starting output-embedding matrix of shape (V, D) as
+                         :class:`Tensor` or :class:`numpy.ndarray`. If None, the Embedding
+                         layer initializes it (seeded if `seed` is set).
+            seed: Optional RNG seed used for **embedding initialization** and for
+                  **negative sampling** during training.
+            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 (required if lr_sched is provided).
+            device: Target device string (e.g., "cuda"); accepted for API parity, ignored by PureML.
         """
+
+        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)
-        self.in_emb  = Embedding(V, D)
-        self.out_emb = Embedding(V, D)
 
+        # Convert warm-starts from np.ndarray → Tensor if needed
+        if isinstance(in_weights, np.ndarray):
+            in_weights = Tensor(in_weights, requires_grad=True)
+        if isinstance(out_weights, np.ndarray):
+            out_weights = Tensor(out_weights, requires_grad=True)
+
+        # embeddings
+        self.in_emb  = Embedding(self.V, self.D, W=in_weights,  seed=seed)
+        self.out_emb = Embedding(self.V, self.D, W=out_weights, seed=seed)
+
+        # seed + RNG for negative sampling
         self.seed = None if seed is None else int(seed)
         self._rng = np.random.default_rng(self.seed)
 
-        self.optim: Optim          = optim(self.parameters, **optim_kwargs)
-        self.lr_sched: LRScheduler = lr_sched(**lr_sched_kwargs)
-        _logger.info("SGNS_PureML init: V=%d D=%d seed=%s", self.V, self.D, self.seed)
+        # API compatibility: PureML is CPU-only
+        self.device = "cpu"
+
+        # optimizer / scheduler
+        self.optim: Optim = optim(self.parameters, **optim_kwargs)
+        self.lr_sched: LRScheduler | None = (
+            lr_sched(optim=self.optim, **lr_sched_kwargs) if lr_sched is not None else None
+        )
+
+        _logger.info(
+            "SGNS_PureML init: V=%d D=%d device=%s seed=%s",
+            self.V, self.D, self.device, self.seed
+        )
 
-    def _sample_neg(self, B: int, K: int, dist: np.ndarray):
-        """Draw negative samples according to the provided unigram distribution."""
-        if dist.ndim != 1 or dist.size != self.V:
-            raise ValueError(f"noise_dist must be 1-D with length {self.V}; got {dist.shape}")
+    def _sample_neg(self, B: int, K: int, dist: np.ndarray) -> np.ndarray:
         return self._rng.choice(self.V, size=(B, K), replace=True, p=dist)
 
-    def predict(self, center: Tensor, pos: Tensor, neg: Tensor) -> Tensor:
-        """Compute positive/negative logits for SGNS."""
-        c      = self.in_emb(center)
-        pos_e  = self.out_emb(pos)
-        neg_e  = self.out_emb(neg)
-        pos_logits = t_sum(c * pos_e, axis=-1)
-        neg_logits = t_sum(c[:, None, :] * neg_e, axis=-1)
-        #                       ^^^
-        # (B,1,D) * (B,K,D) → (B,K,D) → sum D → (B,K)
+    def predict(self, center: Tensor, pos: Tensor, neg: Tensor) -> tuple[Tensor, Tensor]:
+        """Compute positive/negative logits for SGNS.
 
+        Shapes:
+            center: (B,)
+            pos:    (B,)
+            neg:    (B, K)
+        Returns:
+            pos_logits: (B,)
+            neg_logits: (B, K)
+        """
+        c      = self.in_emb(center)      # (B, D)
+        pos_e  = self.out_emb(pos)        # (B, D)
+        neg_e  = self.out_emb(neg)        # (B, K, D)
+
+        pos_logits = t_sum(c * pos_e, axis=-1)                # (B,)
+        neg_logits = t_sum(c[:, None, :] * neg_e, axis=-1)    # (B, K)
         return pos_logits, neg_logits
 
     def fit(self,
@@ -89,47 +136,273 @@ class SGNS_PureML(NN):
             "SGNS_PureML fit: epochs=%d batch=%d negatives=%d shuffle=%s",
             num_epochs, batch_size, num_negative_samples, shuffle_data
         )
-        data = TensorDataset(centers, contexts)
 
+        if noise_dist.ndim != 1 or noise_dist.size != self.V:
+            raise ValueError(f"noise_dist must be 1-D with length {self.V}; got {noise_dist.shape}")
+        dist = np.asarray(noise_dist, dtype=np.float64)
+        if np.any(dist < 0):
+            raise ValueError("noise_dist has negative entries")
+        s = dist.sum()
+        if not np.isfinite(s) or s <= 0:
+            raise ValueError("noise_dist must have positive finite sum")
+        if abs(s - 1.0) > 1e-6:
+            dist = dist / s
+
+        data = TensorDataset(centers, contexts)
         for epoch in range(1, num_epochs + 1):
             epoch_loss = 0.0
             batches = 0
-            for cen, pos in DataLoader(data, batch_size=batch_size, shuffle=shuffle_data):
-                neg = self._sample_neg(batch_size, num_negative_samples, noise_dist)
+            dl_seed = None if self.seed is None else (self.seed + epoch)
+            for cen, pos in DataLoader(data, batch_size=batch_size, shuffle=shuffle_data, seed=dl_seed):
+                B = cen.data.shape[0] if isinstance(cen, Tensor) else len(cen)
 
+                neg_idx_np = self._sample_neg(B, num_negative_samples, dist)
+                neg = Tensor(neg_idx_np, requires_grad=False)
                 x_pos_logits, x_neg_logits = self(cen, pos, neg)
 
-                y_pos = Tensor(np.ones_like(x_pos_logits.data))
-                y_neg = Tensor(np.zeros_like(x_neg_logits.data))
+                y_pos = Tensor(np.ones_like(x_pos_logits.numpy(copy=False)), requires_grad=False)
+                y_neg = Tensor(np.zeros_like(x_neg_logits.numpy(copy=False)), requires_grad=False)
 
-                loss = BCE(y_pos, x_pos_logits, from_logits=True) + BCE(y_neg, x_neg_logits, from_logits=True)
+                K = int(neg.data.shape[1])
+                loss = (
+                    BCE(y_pos, x_pos_logits, from_logits=True)
+                    + Tensor(K)*BCE(y_neg, x_neg_logits, from_logits=True)
+                )
 
                 self.optim.zero_grad()
                 loss.backward()
                 self.optim.step()
-                
-                if lr_step_per_batch:
+
+                if lr_step_per_batch and self.lr_sched is not None:
                     self.lr_sched.step()
 
-                loss_value = float(np.asarray(loss.data).mean())
+                loss_value = float(np.asarray(loss.data))
                 epoch_loss += loss_value
                 batches += 1
                 _logger.debug("Epoch %d batch %d loss=%.6f", epoch, batches, loss_value)
 
-            if not lr_step_per_batch:
+            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 embeddings(self) -> np.ndarray:
-        """Return the input embedding matrix as a NumPy array."""
-        W: Tensor = self.in_emb.parameters[0]
-        return np.asarray(W.data)
+    def in_embeddings(self) -> np.ndarray:
+        W: Tensor = self.in_emb.parameters[0]   # (V, D)
+        if W.shape != (self.V, self.D):
+            raise RuntimeError(
+                "Wrong embedding matrix shape: "
+                "self.in_emb.parameters[0].shape != (V, D)"
+            )
+        arr = W.numpy(copy=True, readonly=True)  # (V, D)
+        _logger.debug("In emb shape: %s", arr.shape)
+        return arr
+
+    @property
+    def out_embeddings(self) -> np.ndarray:
+        W: Tensor = self.out_emb.parameters[0]  # (V, D)
+        if W.shape != (self.V, self.D):
+            raise RuntimeError(
+                "Wrong embedding matrix shape: "
+                "self.out_emb.parameters[0].shape != (V, D)"
+            )
+        arr = W.numpy(copy=True, readonly=True)  # (V, D)
+        _logger.debug("Out emb shape: %s", arr.shape)
+        return arr
+
+    @property
+    def avg_embeddings(self) -> np.ndarray:
+        return 0.5 * (self.in_embeddings + self.out_embeddings)
+
+class SG_PureML(NN):
+    """Plain Skip-Gram (full softmax) in PureML.
+
+    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 = CCE(one_hot(context, V), logits, from_logits=True)
+
+    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:  Tensor | np.ndarray | None = None,
+                 out_weights: Tensor | np.ndarray | None = None,
+                 *,
+                 seed: int | None = None,
+                 optim: Type[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 model (full softmax, **no biases**).
+
+        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 (Tensor or np.ndarray).
+                out_weights: (D, V) or None — copied into W_out if provided (Tensor or np.ndarray).
+
+        Args:
+            V: Vocabulary size (number of nodes/tokens).
+            D: Embedding dimensionality.
+            in_weights: Optional starting matrix for W_in with shape (V, D) as Tensor or np.ndarray.
+            out_weights: Optional starting matrix for W_out with shape (D, V) as Tensor or np.ndarray.
+                         (Note the asymmetry with SGNS; use `.T` if converting from (V, D).)
+            seed: Optional RNG seed (used for layer initialization).
+            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 (required if lr_sched is provided).
+            device: Device string (e.g., "cuda"). Accepted for parity, ignored by PureML (CPU-only).
+        """
+
+        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)
+
+        # Convert warm-starts from np.ndarray → Tensor if needed
+        if isinstance(in_weights, np.ndarray):
+            in_weights = Tensor(in_weights, requires_grad=True)
+        if isinstance(out_weights, np.ndarray):
+            out_weights = Tensor(out_weights, requires_grad=True)
+
+        # input/output “embedding” projections
+        self.in_emb  = Affine(self.V, self.D, W=in_weights,  bias=False, seed=seed)
+        self.out_emb = Affine(self.D, self.V, W=out_weights, bias=False, seed=seed)
+
+        self.seed = None if seed is None else int(seed)
+        self.device = "cpu"  # API parity
+
+        # optimizer / scheduler
+        self.optim: Optim = optim(self.parameters, **optim_kwargs)
+        self.lr_sched: LRScheduler | None = (
+            lr_sched(optim=self.optim, **lr_sched_kwargs) if lr_sched is not None else None
+        )
+
+        _logger.info(
+            "SG_PureML init: V=%d D=%d device=%s seed=%s",
+            self.V, self.D, self.device, self.seed
+        )
+
+    def predict(self, center: Tensor) -> Tensor:
+        """Return vocabulary logits for each center index.
+
+        Args:
+            center: Tensor of center indices with shape `(B,)` and integer dtype.
+
+        Returns:
+            Tensor: Logits over the vocabulary with shape `(B, V)`.
+        """
+        c = one_hot(dims=self.V, label=center)  # (B, V)
+        y = self.in_emb(c)                      # (B, D)
+        z = self.out_emb(y)                     # (B, V)
+        return z
+
+    def fit(self,
+            centers: np.ndarray,
+            contexts: np.ndarray,
+            num_epochs: int,
+            batch_size: int,
+            shuffle_data: bool,
+            lr_step_per_batch: bool,
+            **_ignore):
+        """Train Skip-Gram with full softmax on center/context pairs.
+
+        Args:
+            centers: Array of center indices, shape `(N,)`, dtype integer in `[0, V)`.
+            contexts: Array of context (target) indices, shape `(N,)`, dtype integer.
+            num_epochs: Number of passes over the dataset.
+            batch_size: Mini-batch size.
+            shuffle_data: Whether to shuffle pairs each epoch.
+            lr_step_per_batch: If True, call `lr_sched.step()` after every batch
+                (when a scheduler is provided). If False, step once per epoch.
+            **_ignore: Ignored kwargs for API compatibility with SGNS.
+
+        Optimization:
+            Uses `CCE(one_hot(context), logits, from_logits=True)` where
+            `logits = predict(center)`. Scheduler stepping obeys `lr_step_per_batch`.
+        """
+        _logger.info(
+            "SG_PureML fit: epochs=%d batch=%d shuffle=%s",
+            num_epochs, batch_size, shuffle_data
+        )
+        data = TensorDataset(centers, contexts)
+
+        for epoch in range(1, num_epochs + 1):
+            epoch_loss = 0.0
+            batches = 0
+            dl_seed = None if self.seed is None else (self.seed + epoch)
+            for cen, ctx in DataLoader(data, batch_size=batch_size, shuffle=shuffle_data, seed=dl_seed):
+                logits = self(cen)                          # (B, V)
+                y = one_hot(dims=self.V, label=ctx)         # (B, V)
+                loss = CCE(y, logits, from_logits=True)     # scalar
+
+                self.optim.zero_grad()
+                loss.backward()
+                self.optim.step()
+
+                if lr_step_per_batch and self.lr_sched is not None:
+                    self.lr_sched.step()
+
+                loss_value = float(np.asarray(loss.data))
+                epoch_loss += loss_value
+                batches += 1
+                _logger.debug("Epoch %d batch %d loss=%.6f", epoch, batches, loss_value)
+
+            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:
+        """Input embeddings matrix `W_in` as `(V, D)` (copy, read-only)."""
+        W = self.in_emb.parameters[0]              # (V, D)
+        if W.shape != (self.V, self.D):
+            raise RuntimeError(
+                "Wrong embedding matrix shape: "
+                "self.in_emb.parameters[0].shape != (V, D)"
+            )
+        arr = W.numpy(copy=True, readonly=True)    # (V, D)
+        _logger.debug("In emb shape: %s", arr.shape)
+        return arr
+
+    @property
+    def out_embeddings(self) -> np.ndarray:
+        """Output embeddings matrix `W_outᵀ` as `(V, D)` (copy, read-only).
+        (`out_emb.parameters[0]` is `(D, V)`, so we transpose.)"""
+        W = self.out_emb.parameters[0]             # (D, V)
+        if W.shape != (self.D, self.V):
+            raise RuntimeError(
+                "Wrong embedding matrix shape: "
+                "self.out_emb.parameters[0].shape != (D, V)"
+            )
+        arr = W.numpy(copy=True, readonly=True).T  # (V, D)
+        _logger.debug("Out emb shape: %s", arr.shape)
+        return arr
+    
+    @property
+    def avg_embeddings(self) -> np.ndarray:
+        """Elementwise average of input/output embeddings, shape `(V, D)`."""
+        return 0.5 * (self.in_embeddings + self.out_embeddings)  # (V, D)
 
 
-__all__ = ["SGNS_PureML"]
+__all__ = ["SGNS_PureML", "SG_PureML"]
 
 if __name__ == "__main__":
     pass