sawnergy 1.0.5__py3-none-any.whl → 1.0.7__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.training_utils import TensorDataset, DataLoader, one_hot
 from pureml.base import NN
 
 # built-in
@@ -34,45 +34,73 @@ class SGNS_PureML(NN):
                 seed: int | None = None,
                 optim: Type[Optim],
                 optim_kwargs: dict,
-                lr_sched: Type[LRScheduler],
-                lr_sched_kwargs: dict):
+                lr_sched: Type[LRScheduler] | None = None,
+                lr_sched_kwargs: dict | None = None,
+                device: str | None = None):
         """
         Args:
             V: Vocabulary size (number of nodes).
             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.
+            optim: Optimizer class to instantiate.
+            optim_kwargs: Keyword arguments for the optimizer (required).
+            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")
+        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)
 
+        # embeddings
+        self.in_emb  = Embedding(self.V, self.D)
+        self.out_emb = Embedding(self.V, self.D)
+
+        # seed + RNG for negative sampling
         self.seed = None if seed is None else int(seed)
         self._rng = np.random.default_rng(self.seed)
+        if self.seed is not None:
+            # optional: also set global NumPy seed for any non-RNG paths
+            np.random.seed(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"
 
-    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}")
+        # 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) -> 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 +117,254 @@ 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)
+                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)
+                    + 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)"
+            )
+        return W.numpy(copy=True, readonly=True)
+
+    @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)"
+            )
+        return W.numpy(copy=True, readonly=True)
+
+    @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],
+                 optim_kwargs: dict,
+                 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 (e.g., `Adam`, `SGD`).
+            optim_kwargs: Keyword arguments passed to the optimizer constructor.
+            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).
+
+        Notes:
+            The encoder/decoder are implemented as:
+              • `in_emb = Affine(V, D)` (acts on a one-hot center index)
+              • `out_emb = Affine(D, V)`
+            so forward pass produces vocabulary-sized logits.
+        """
+        if optim_kwargs is None:
+            raise ValueError("optim_kwargs must be provided")
+        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)
+
+        # 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(
+            "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)"
+            )
+        return W.numpy(copy=True, readonly=True)   # (V, D)
+
+    @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)"
+            )
+        return W.numpy(copy=True, readonly=True).T # (V, D)
+    
+    @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

@@ -42,31 +42,31 @@ class SGNS_Torch:
             optim: Optimizer class to instantiate.
             optim_kwargs: Keyword arguments for the optimizer.
             lr_sched: Optional learning-rate scheduler class.
-            lr_sched_kwargs: Keyword arguments for the scheduler.
+            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")
         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 +112,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 +149,7 @@ 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()
@@ -163,6 +172,131 @@ class SGNS_Torch:
     def embeddings(self) -> np.ndarray:
         """Return the input embedding matrix as a NumPy array."""
         return self.in_emb.weight.detach().cpu().numpy()
+    
+    @property
+    def in_embeddings(self) -> np.ndarray:
+        return self.in_emb.weight.detach().cpu().numpy()
+
+    @property
+    def out_embeddings(self) -> np.ndarray:
+        return self.out_emb.weight.detach().cpu().numpy()
+
+    @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:
+
+    def __init__(self,
+                 V: int,
+                 D: int,
+                 *,
+                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):
+
+        if optim_kwargs is None:
+            raise ValueError("optim_kwargs must be provided")
+        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()
+                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:
+        return self.in_emb.weight.detach().T.cpu().numpy()
+
+    @property
+    def out_embeddings(self) -> np.ndarray:
+        return self.out_emb.weight.detach().cpu().numpy()
+
+    @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 +305,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"
 ]