sawnergy 1.0.6__py3-none-any.whl → 1.0.8__py3-none-any.whl

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
@@ -32,8 +32,8 @@ class SGNS_PureML(NN):
                 D: int,
                 *,
                 seed: int | None = None,
-                optim: Type[Optim],
-                optim_kwargs: dict,
+                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):
@@ -42,15 +42,15 @@ class SGNS_PureML(NN):
             V: Vocabulary size (number of nodes).
             D: Embedding dimensionality.
             seed: Optional RNG seed for negative sampling.
-            optim: Optimizer class to instantiate.
-            optim_kwargs: Keyword arguments for the optimizer (required).
+            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.
         """
 
-        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")
 
@@ -82,9 +82,6 @@ class SGNS_PureML(NN):
         )
 
     def _sample_neg(self, B: int, K: int, dist: np.ndarray) -> 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}")
         return self._rng.choice(self.V, size=(B, K), replace=True, p=dist)
 
     def predict(self, center: Tensor, pos: Tensor, neg: Tensor) -> tuple[Tensor, Tensor]:
@@ -120,8 +117,19 @@ 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
@@ -129,16 +137,17 @@ class SGNS_PureML(NN):
             for cen, pos in DataLoader(data, batch_size=batch_size, shuffle=shuffle_data):
                 B = cen.data.shape[0] if isinstance(cen, Tensor) else len(cen)
 
-                neg_idx_np = self._sample_neg(B, num_negative_samples, noise_dist)
+                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)
 
+                K = int(neg.data.shape[1])
                 loss = (
                     BCE(y_pos, x_pos_logits, from_logits=True)
-                    + BCE(y_neg, x_neg_logits, from_logits=True)
+                    + Tensor(K)*BCE(y_neg, x_neg_logits, from_logits=True)
                 )
 
                 self.optim.zero_grad()
@@ -160,13 +169,200 @@ class SGNS_PureML(NN):
             _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 (V, D)."""
-        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.
+
+    Trains two affine layers to emulate the classic Skip-Gram objective with a
+    **full** softmax over the vocabulary (no negative sampling):
+
+        x = one_hot(center, V)                 # (B, V)
+        y = x @ W_in + b_in                    # (B, D)
+        logits = y @ W_out + b_out             # (B, V)
+        loss = CCE(one_hot(context, V), logits, from_logits=True)
+
+    The learnable “input” embeddings are the rows of `W_in` (shape `(V, D)`), and
+    the “output” embeddings are the rows of `W_outᵀ` (also `(V, D)`).
+    """
+
+    def __init__(self,
+                V: int,
+                D: int,
+                *,
+                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).
+
+        Args:
+            V: Vocabulary size (number of nodes/tokens).
+            D: Embedding dimensionality.
+            seed: Optional RNG seed (kept for API parity; not used in layer init).
+            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)
+
+        # input/output “embedding” projections
+        self.in_emb  = Affine(self.V, self.D)
+        self.out_emb = Affine(self.D, self.V)
+
+        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
+
+            for cen, ctx in DataLoader(data, batch_size=batch_size, shuffle=shuffle_data):
+                logits = self(cen)                          # (B, V)
+                y = one_hot(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

sawnergy/embedding/SGNS_torch.py

@@ -10,6 +10,7 @@ from torch.optim.lr_scheduler import LRScheduler
 # built-in
 import logging
 from typing import Type
+import warnings
 
 # *----------------------------------------------------*
 #                        GLOBALS
@@ -22,51 +23,84 @@ _logger = logging.getLogger(__name__)
 # *----------------------------------------------------*
 
 class SGNS_Torch:
-    """PyTorch implementation of Skip-Gram with Negative Sampling."""
+    """PyTorch implementation of Skip-Gram with Negative Sampling.
+
+    DEPRECATED (temporary): This class currently produces noisy embeddings in
+    practice and is deprecated until further notice. The issue likely stems from
+    weight initialization, although the root cause has not yet been determined.
+
+    Prefer one of the following alternatives:
+      • Plain PyTorch Skip-Gram (full softmax): `SG_Torch`
+      • PureML-based implementations: `SGNS_PureML` or `SG_PureML` (if available)
+
+    This API may change or be removed once the root cause is resolved.
+    """
 
     def __init__(self,
-                 V: int,
-                 D: int,
-                 *,
+                V: int,
+                D: int,
+                *,
                 seed: int | None = None,
-                optim: Type[Optimizer],
-                optim_kwargs: dict,
+                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.
+
+        DEPRECATION WARNING:
+            This implementation is temporarily deprecated for producing noisy
+            embeddings. The issue likely stems from weight initialization, though
+            the exact root cause has not been conclusively determined. Please use
+            `SG_Torch` (plain Skip-Gram with full softmax) or the PureML-based
+            `SGNS_PureML` / `SG_PureML` models instead.
+
         Args:
             V: Vocabulary size (number of nodes).
             D: Embedding dimensionality.
             seed: Optional RNG seed for PyTorch.
-            optim: Optimizer class to instantiate.
-            optim_kwargs: Keyword arguments for the optimizer.
+            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")
+
+        # --- runtime deprecation notice ---
+        warnings.warn(
+            "SGNS_Torch is temporarily deprecated: it currently produces noisy "
+            "embeddings (likely due to weight initialization). Use SG_Torch "
+            "(plain Skip-Gram, full softmax) or the PureML-based SG/SGNS classes.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        _logger.warning(
+            "DEPRECATED: SGNS_Torch currently produces noisy embeddings "
+            "(likely weight initialization). Prefer SG_Torch or PureML SG/SGNS."
+        )
+        # ----------------------------------
+
+        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")
-        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)
+        # two embeddings as in/out matrices
+        self.in_emb  = nn.Embedding(self.V, self.D)
+        self.out_emb = nn.Embedding(self.V, self.D)
 
+        resolved_device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = torch.device(resolved_device)
         if seed is not None:
             torch.manual_seed(int(seed))
             np.random.seed(int(seed))
             if self.device.type == "cuda":
                 torch.cuda.manual_seed_all(int(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.to(self.device)
-
+        _logger.info("SGNS_Torch init: V=%d D=%d device=%s seed=%s", self.V, self.D, self.device, 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
 
@@ -112,6 +146,15 @@ class SGNS_Torch:
         idx = np.arange(N)
 
         noise_probs = torch.as_tensor(noise_dist, dtype=torch.float32, device=self.device)
+        # require normalized, non-negative distribution
+        if (not torch.isfinite(noise_probs).all()
+            or (noise_probs < 0).any()
+            or abs(float(noise_probs.sum().item()) - 1.0) > 1e-6):
+            raise ValueError(
+                "noise_dist must be non-negative, finite, and sum to 1.0 "
+                f"(got sum={float(noise_probs.sum().item()):.6f}, "
+                f"min={float(noise_probs.min().item()):.6f})"
+            )
 
         for epoch in range(1, int(num_epochs) + 1):
             epoch_loss = 0.0
@@ -140,7 +183,140 @@ class SGNS_Torch:
                 y_neg = torch.zeros_like(neg_logits)
                 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()
+        _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()
+        _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."""
+
+    def __init__(self,
+                V: int,
+                D: int,
+                *,
+                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) model in PyTorch.
+
+        Args:
+            V: Vocabulary size (number of nodes/tokens).
+            D: Embedding dimensionality.
+            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 linear layers acting on one-hot centers:
+            • ``in_emb = nn.Linear(V, D)``
+            • ``out_emb = nn.Linear(D, V)``
+            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)
+
+        self.in_emb  = nn.Linear(self.V, self.D)
+        self.out_emb = nn.Linear(self.D, self.V)
+
+        resolved_device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = torch.device(resolved_device)
+        if seed is not None:
+            torch.manual_seed(int(seed))
+            np.random.seed(int(seed))
+            if self.device.type == "cuda":
+                torch.cuda.manual_seed_all(int(seed))
+        self.to(self.device)
+        _logger.info("SG_Torch init: V=%d D=%d device=%s seed=%s", self.V, self.D, self.device, 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:
+                np.random.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()
@@ -158,11 +334,22 @@ class SGNS_Torch:
 
             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().T.cpu().numpy()
+        _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()
+        _logger.debug("Out emb shape: %s", W.shape)
+        return W
 
     @property
-    def embeddings(self) -> np.ndarray:
-        """Return the input embedding matrix as a NumPy array."""
-        return self.in_emb.weight.detach().cpu().numpy()
+    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 +358,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 .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,24 @@ 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",
+    "Visualizer",
     "SGNS_PureML",
     "SGNS_Torch",
+    "SG_PureML",
+    "SG_Torch"
 ]