kaiwu-torch-plugin 0.2.0__py3-none-any.whl

kaiwu/torch_plugin/__init__.py

@@ -0,0 +1,19 @@
+# -*- coding: utf-8 -*-
+"""Kaiwu PyTorch plugin public API."""
+
+from .dbn import UnsupervisedDBN
+from .full_boltzmann_machine import BoltzmannMachine
+from .qdiffusion import QDiffusion, QDiffusionConfig
+from .qvae import QVAE
+from .restricted_boltzmann_machine import RestrictedBoltzmannMachine
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "RestrictedBoltzmannMachine",
+    "BoltzmannMachine",
+    "QVAE",
+    "UnsupervisedDBN",
+    "QDiffusion",
+    "QDiffusionConfig",
+]

kaiwu/torch_plugin/_qdiffusion_sampling.py

@@ -0,0 +1,158 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+# SPDX-License-Identifier: Apache-2.0
+
+"""Private sampling helpers shared by the public QDiffusion model."""
+
+from __future__ import annotations
+
+import torch
+
+
+# Skeptical-remasking helpers.
+
+def topk_masking(
+    scores: torch.Tensor,
+    cutoff_len: torch.Tensor,
+    stochastic: bool = False,
+    temp: float = 1.0,
+) -> torch.Tensor:
+    """Selects the lowest-score positions used by skeptical remasking.
+
+    Args:
+        scores: Per-position scores used to rank editable token positions.
+        cutoff_len: Per-sample cutoff lengths that determine how many
+            positions remain masked.
+        stochastic: Whether to perturb the ranking with Gumbel noise.
+        temp: Noise temperature applied when ``stochastic`` is enabled.
+
+    Returns:
+        torch.Tensor: A boolean mask where ``True`` marks positions below the per-sample
+        cutoff.
+    """
+    if stochastic:
+        gumbel_noise = -torch.log(-torch.log(torch.rand_like(scores) + 1e-8) + 1e-8)
+        ranked_scores = scores + temp * gumbel_noise
+    else:
+        ranked_scores = scores
+    sorted_scores = ranked_scores.sort(-1)[0]
+    cutoff = sorted_scores.gather(dim=-1, index=cutoff_len)
+    return ranked_scores < cutoff
+
+
+# Categorical sampling helpers.
+
+def sample_from_categorical(
+    logits: torch.Tensor,
+    temperature: float = 1.0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Samples tokens from categorical logits.
+
+    Args:
+        logits: Unnormalized categorical logits.
+        temperature: Sampling temperature. A falsy value switches to greedy
+            argmax decoding.
+
+    Returns:
+        tuple[torch.Tensor, torch.Tensor]: A tuple ``(tokens, scores)``
+        where ``tokens`` contains sampled token ids and ``scores`` contains
+        the associated log-probability-style scores.
+    """
+    if temperature:
+        dist = torch.distributions.Categorical(logits=logits.div(temperature))
+        tokens = dist.sample()
+        scores = dist.log_prob(tokens)
+    else:
+        scores, tokens = logits.log_softmax(dim=-1).max(dim=-1)
+    return tokens, scores
+
+
+def stochastic_sample_from_categorical(
+    logits: torch.Tensor,
+    temperature: float = 1.0,
+    noise_scale: float = 1.0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Applies Gumbel noise before categorical sampling.
+
+    Args:
+        logits: Unnormalized categorical logits.
+        temperature: Sampling temperature forwarded to
+            :func:`sample_from_categorical`.
+        noise_scale: Multiplicative scale for the sampled Gumbel noise.
+
+    Returns:
+        tuple[torch.Tensor, torch.Tensor]: A tuple ``(tokens, scores)``
+        sampled from the perturbed categorical distribution.
+    """
+    gumbel_noise = -torch.log(-torch.log(torch.rand_like(logits) + 1e-8) + 1e-8)
+    noisy_logits = logits + noise_scale * gumbel_noise
+    return sample_from_categorical(noisy_logits, temperature)
+
+
+def stochastic_sample_from_categorical_n(
+    logits: torch.Tensor,
+    temperature: float = 1.0,
+    noise_scale: float = 1.0,
+    n: int = 1,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Samples multiple noisy categorical candidates.
+
+    Args:
+        logits: Unnormalized categorical logits shaped as
+            ``[batch, seq_len, vocab]`` or similar.
+        temperature: Sampling temperature forwarded to
+            :func:`sample_from_categorical`.
+        noise_scale: Multiplicative scale for the sampled Gumbel noise.
+        n: Number of independent noisy candidate sets to draw.
+
+    Returns:
+        tuple[torch.Tensor, torch.Tensor]: A tuple ``(tokens, scores)``
+        whose leading dimension indexes the sampled candidate set.
+    """
+    expanded_logits = logits.unsqueeze(0).expand(n, *logits.shape)
+    gumbel_noise = -torch.log(
+        -torch.log(torch.rand_like(expanded_logits) + 1e-8) + 1e-8
+    )
+    noisy_logits = expanded_logits + noise_scale * gumbel_noise
+    return sample_from_categorical(noisy_logits, temperature)
+
+
+# Logit filtering helpers.
+
+def top_k_top_p_filtering(
+    logits: torch.Tensor,
+    top_k: int = 0,
+    top_p: float = 0.95,
+    filter_value: float = -float("Inf"),
+) -> torch.Tensor:
+    """Applies top-k and/or nucleus filtering to logits.
+
+    Args:
+        logits: Unnormalized categorical logits.
+        top_k: Number of highest-logit entries to retain per row. A value of
+            ``0`` disables top-k filtering.
+        top_p: Nucleus-filtering threshold on cumulative probability mass.
+        filter_value: Replacement value written into filtered logit entries.
+
+    Returns:
+        torch.Tensor: A tensor with the same shape as ``logits`` where filtered entries are
+        replaced by ``filter_value``.
+    """
+    original_shape = logits.shape
+    flat_logits = logits.reshape(-1, original_shape[-1])
+    assert flat_logits.dim() == 2
+
+    top_k = min(top_k, flat_logits.size(-1))
+    if top_k > 0:
+        indices_to_remove = (
+            flat_logits < torch.topk(flat_logits, top_k, dim=1)[0][..., -1, None]
+        )
+        flat_logits[indices_to_remove] = filter_value
+
+    sorted_logits, sorted_indices = torch.sort(flat_logits, descending=True)
+    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+    sorted_indices_to_remove = cumulative_probs > top_p
+    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+    sorted_indices_to_remove[..., 0] = 0
+    sorted_logits[sorted_indices_to_remove] = filter_value
+    restored_logits = torch.gather(sorted_logits, 1, sorted_indices.argsort(-1))
+    return restored_logits.reshape(original_shape)

kaiwu/torch_plugin/abstract_boltzmann_machine.py

@@ -0,0 +1,102 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2022-2025 Beijing QBoson Quantum Technology Co., Ltd.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+"""Abstract base class for Boltzmann Machines."""
+import torch
+
+
+class AbstractBoltzmannMachine(torch.nn.Module):
+    """Abstract base class for Boltzmann Machines.
+
+    Args:
+        device (torch.device, optional): Device for tensor construction.
+    """
+
+    def __init__(self, device=None) -> None:
+        super().__init__()
+        if device is None:
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        else:
+            self.device = device
+
+    def to(self, device=..., dtype=..., non_blocking=...):
+        """Moves the model to the specified device.
+
+        Args:
+            device: Target device.
+            dtype: Target data type.
+            non_blocking: Whether the operation should be non-blocking.
+
+        Returns:
+            AbstractBoltzmannMachine: The model on the target device.
+        """
+        self.device = device
+        return super().to(device)
+
+    def forward(self, s_all: torch.Tensor) -> torch.Tensor:
+        """Computes the Hamiltonian.
+
+        Args:
+            s_all (torch.Tensor): Input tensor.
+
+        Returns:
+            torch.Tensor: Hamiltonian.
+        """
+
+    def get_ising_matrix(self):
+        """Converts the model to Ising format.
+
+        Returns:
+            torch.Tensor: Ising matrix.
+        """
+        return self._to_ising_matrix()
+
+    def _to_ising_matrix(self):
+        """Converts the model to Ising format.
+
+        Returns:
+            torch.Tensor: Ising matrix.
+
+        Raises:
+            NotImplementedError: If not implemented in subclass.
+        """
+        raise NotImplementedError("Subclasses must implement _ising method")
+
+    def objective(
+        self,
+        s_positive: torch.Tensor,
+        s_negative: torch.Tensor,
+    ) -> torch.Tensor:
+        """Objective function whose gradient is equivalent to the gradient of
+        negative log-likelihood.
+
+        Args:
+            s_positive (torch.Tensor): Tensor of observed spins (data), shape (b1, N),
+                            where b1 is batch size and N is the number of variables.
+            s_negative (torch.Tensor): Tensor of spins sampled from the model, shape (b2, N),
+                            where b2 is batch size and N is the number of variables.
+
+        Returns:
+            torch.Tensor: Scalar difference between data and model average energy.
+        """
+        return self(s_positive).mean() - self(s_negative).mean()
+
+    def sample(self, sampler) -> torch.Tensor:
+        """Samples from the Boltzmann Machine.
+
+        Args:
+            sampler (kaiwu.core.OptimizerBase): Optimizer used for sampling from the model.
+                The sampler can be kaiwuSDK's CIM or other solvers.
+
+        Returns:
+            torch.Tensor: Spins sampled from the model.
+        """
+        ising_mat = self.get_ising_matrix()
+        solution = sampler.solve(ising_mat)
+        solution = (solution[:, :-1] * solution[:, [-1]] + 1) / 2
+        solution = torch.FloatTensor(solution)
+        solution = solution.to(self.device)
+        return solution

kaiwu/torch_plugin/dbn.py

@@ -0,0 +1,220 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2022-2025 Beijing QBoson Quantum Technology Co., Ltd.
+#
+# SPDX-License-Identifier: Apache-2.0
+"""Deep Belief Network (DBN) model.
+
+This module contains the DBN class and functions for training the DBN+model
+or only the model. Training the DBN+model will save the likelihood values
+and prediction accuracy during the training process.
+"""
+import numpy as np
+
+
+import torch
+from torch import nn
+
+from .restricted_boltzmann_machine import RestrictedBoltzmannMachine
+
+
+# =================== Unsupervised DBN General Model =====================
+class UnsupervisedDBN(nn.Module):
+    """A general unsupervised Deep Belief Network (DBN) architecture.
+
+    This model is a stack of Restricted Boltzmann Machines (RBMs).
+
+    Args:
+        hidden_layers_structure (list, optional): A list of integers
+            representing the number of hidden units in each layer.
+            Defaults to [100, 100].
+    """
+
+    def __init__(self, hidden_layers_structure=None):
+        super().__init__()
+        self.hidden_layers_structure = (
+            hidden_layers_structure
+            if hidden_layers_structure is not None
+            else [100, 100]
+        )
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.rbm_layers = None
+        self.input_dim = None
+        self._is_trained = False
+
+    def create_rbm_layer(self, input_dim):
+        """Creates the layers of RBMs for the DBN.
+
+        Args:
+            input_dim (int): The dimension of the input data (number of visible units).
+
+        Returns:
+            UnsupervisedDBN: The instance itself with the RBM layers created.
+        """
+        self.input_dim = input_dim
+        self.rbm_layers = nn.ModuleList()
+
+        current_dim = input_dim
+        for n_hidden in self.hidden_layers_structure:
+            rbm = RestrictedBoltzmannMachine(
+                num_visible=current_dim,  # Number of visible units (feature dimension)
+                num_hidden=n_hidden,  # Number of hidden units
+            ).to(
+                self.device
+            )  # Move model to specified device (CPU/GPU)
+            self.rbm_layers.append(rbm)
+            current_dim = n_hidden
+
+        self._is_trained = False
+        return self
+
+    def forward(self, data_in):
+        """Performs a forward pass to transform the input data.
+
+        Args:
+            data_in (numpy.ndarray): The input data.
+
+        Returns:
+            numpy.ndarray: The transformed data after passing through all RBM layers.
+
+        Raises:
+            ValueError: If the model has not been built or trained yet.
+        """
+        if self.rbm_layers is None:
+            raise ValueError("Model not built yet. Call create_rbm_layer first.")
+        if not self._is_trained:
+            raise ValueError(
+                "Model not trained yet. Call mark_as_trained() after training."
+            )
+
+        data_in = data_in.astype(np.float32)
+        for rbm in self.rbm_layers:
+            with torch.no_grad():
+                hidden_output = rbm.get_hidden(
+                    torch.FloatTensor(data_in).to(self.device)
+                )
+                data_in = (
+                    hidden_output[:, rbm.num_visible :].cpu().numpy()
+                )  # Extract only the hidden part
+        return data_in
+
+    def transform(self, data_in):
+        """An sklearn-compatible transform method.
+
+        Args:
+            data_in (numpy.ndarray): The input data.
+
+        Returns:
+            numpy.ndarray: The transformed data.
+        """
+        return self.forward(data_in)
+
+    def reconstruct(self, data_in, layer_index=0):
+        """Reconstructs the input from a specified RBM layer.
+
+        Args:
+            data_in (numpy.ndarray): The input data to be reconstructed.
+
+            layer_index (int, optional): The index of the RBM layer to use for reconstruction.
+                Defaults to 0.
+
+        Returns:
+            numpy.ndarray: The reconstructed data.
+
+        Raises:
+            ValueError: If the model has no RBM layers or the layer index is out of range.
+        """
+        if self.rbm_layers is None or len(self.rbm_layers) == 0:
+            raise ValueError("No RBM layers found. Please fit the model first.")
+
+        if layer_index >= len(self.rbm_layers):
+            raise ValueError(f"Layer index {layer_index} out of range.")
+
+        rbm = self.rbm_layers[layer_index]
+        return self.reconstruct_with_rbm(rbm, data_in, self.device)
+
+    def mark_as_trained(self):
+        """Marks the model as trained.
+
+        Returns:
+            UnsupervisedDBN: The instance itself.
+        """
+        self._is_trained = True
+        return self
+
+    def get_rbm_layer(self, index):
+        """Gets the RBM layer at the specified index.
+
+        Args:
+            index (int): The index of the RBM layer.
+
+        Returns:
+            RestrictedBoltzmannMachine or None: The RBM layer if found, otherwise None.
+        """
+        if index < len(self.rbm_layers):
+            return self.rbm_layers[index]
+        return None
+
+    @staticmethod
+    def reconstruct_with_rbm(rbm, data_in, device=None):
+        """Reconstructs data using a single RBM.
+
+        Args:
+            rbm (RestrictedBoltzmannMachine): The trained RBM model.
+
+            data_in (numpy.ndarray): The input data.
+
+            device (torch.device, optional): The device to perform computation on.
+                If None, uses the RBM's device. Defaults to None.
+
+        Returns:
+            tuple[numpy.ndarray, numpy.ndarray]: A tuple containing:
+                - The reconstructed visible layer data.
+                - The reconstruction error for each sample.
+        """
+        if device is None:
+            device = rbm.device
+
+        # Convert to PyTorch tensor
+        data_in = torch.FloatTensor(data_in).to(device)
+
+        with torch.no_grad():
+            # Get hidden representation using RBM's get_hidden
+            hidden_act = rbm.get_hidden(data_in)
+            hidden_part = hidden_act[
+                :, rbm.num_visible :
+            ]  # Extract only the hidden part
+
+            # Reconstruct visible layer (using transposed weights)
+            visible_recon = torch.sigmoid(
+                torch.matmul(hidden_part, rbm.quadratic_coef.t())
+                + rbm.linear_bias[: rbm.num_visible]
+            )
+
+            # Calculate reconstruction error
+            recon_errors = (
+                torch.mean((data_in - visible_recon) ** 2, dim=1).cpu().numpy()
+            )
+
+        return visible_recon.cpu().numpy(), recon_errors
+
+    @property
+    def num_layers(self):
+        """Returns the number of RBM layers.
+
+        Returns:
+            int: The number of layers.
+        """
+        return len(self.rbm_layers)
+
+    @property
+    def output_dim(self):
+        """Returns the output dimension of the DBN.
+
+        Returns:
+            int: The dimension of the final hidden layer.
+        """
+        if len(self.rbm_layers) > 0:
+            return self.rbm_layers[-1].num_hidden
+        return self.input_dim