parameters

        beta: float = 0.1,
        pad_id: int = None,
    ):
        """
        Initializes the DPO model.

        Args:
            model (TransformerWrapper): The policy model to be optimized.
            beta (float, optional): The beta parameter for the divergence metric. Defaults to 0.1.
            pad_id (int, optional): The padding token ID. Defaults to None.
        """
        super().__init__()
        self.policy_model = model

        self.ref_model = deepcopy(model)
        freeze_all_layers_(self.ref_model)

        self.beta = beta
        self.pad_id = pad_id

    def parameters(self) -> Iterator[torch.nn.Parameter]:
        """
        Returns an iterator over the model parameters.

        Returns:
            Iterator[torch.nn.Parameter]: An iterator over the model parameters.
        """
        return self.policy_model.parameters()

    def forward(
        self,
        preferred_seq: torch.Tensor,
        unpreferred_seq: torch.Tensor,
        *,
        prompt_mask: torch.Tensor,
        preferred_seq_mask: torch.Tensor = None,
        unpreferred_seq_mask: torch.Tensor = None,
    ) -> torch.Tensor:
        """
        Computes the DPO loss for the given sequences and masks.

        Args:
            preferred_seq (torch.Tensor): The preferred sequence tensor.
            unpreferred_seq (torch.Tensor): The unpreferred sequence tensor.
            prompt_mask (torch.Tensor): The prompt mask tensor.
            preferred_seq_mask (torch.Tensor, optional): The mask for the preferred sequence. Defaults to None.
            unpreferred_seq_mask (torch.Tensor, optional): The mask for the unpreferred sequence. Defaults to None.


forward

            pad_id (int, optional): The padding token ID. Defaults to None.
        """
        super().__init__()
        self.policy_model = model

        self.ref_model = deepcopy(model)
        freeze_all_layers_(self.ref_model)

        self.beta = beta
        self.pad_id = pad_id

    def parameters(self) -> Iterator[torch.nn.Parameter]:
        """
        Returns an iterator over the model parameters.

        Returns:
            Iterator[torch.nn.Parameter]: An iterator over the model parameters.
        """
        return self.policy_model.parameters()

    def forward(
        self,
        preferred_seq: torch.Tensor,
        unpreferred_seq: torch.Tensor,
        *,
        prompt_mask: torch.Tensor,
        preferred_seq_mask: torch.Tensor = None,
        unpreferred_seq_mask: torch.Tensor = None,
    ) -> torch.Tensor:
        """
        Computes the DPO loss for the given sequences and masks.

        Args:
            preferred_seq (torch.Tensor): The preferred sequence tensor.
            unpreferred_seq (torch.Tensor): The unpreferred sequence tensor.
            prompt_mask (torch.Tensor): The prompt mask tensor.
            preferred_seq_mask (torch.Tensor, optional): The mask for the preferred sequence. Defaults to None.
            unpreferred_seq_mask (torch.Tensor, optional): The mask for the unpreferred sequence. Defaults to None.

        Returns:
            torch.Tensor: The computed DPO loss tensor.
        """
        assert preferred_seq.ndim == 2
        assert preferred_seq.shape == unpreferred_seq.shape

        if exists(self.pad_id):
            if not exists(preferred_seq_mask):
                preferred_seq_mask = preferred_seq != self.pad_id

            if not exists(unpreferred_seq_mask):
                unpreferred_seq_mask = unpreferred_seq != self.pad_id

        """
        Following Appendix B in https://arxiv.org/abs/2305.18290
        """

        with torch.no_grad():
            self.ref_model.eval()
            ref_preferred_logprob = log_prob_from_model_and_seq(
                self.ref_model, preferred_seq
            )
            ref_unpreferred_logprob = log_prob_from_model_and_seq(
                self.ref_model, unpreferred_seq
            )

        policy_preferred_logprob = log_prob_from_model_and_seq(
            self.policy_model, preferred_seq
        )
        policy_unpreferred_logprob = log_prob_from_model_and_seq(
            self.policy_model, unpreferred_seq
        )

        # masked mean of log probs

        preferred_seq_mask = maybe_and_mask(
            ~prompt_mask, preferred_seq_mask
        )
        unpreferred_seq_mask = maybe_and_mask(
            ~prompt_mask, unpreferred_seq_mask
        )

        ref_preferred_logprob, policy_preferred_logprob = map(
            lambda t: masked_mean(t, preferred_seq_mask),
            (ref_preferred_logprob, policy_preferred_logprob),
        )
        ref_unpreferred_logprob, policy_unpreferred_logprob = map(
            lambda t: masked_mean(t, unpreferred_seq_mask),
            (ref_unpreferred_logprob, policy_unpreferred_logprob),
        )

        # main dpo formula

        policy_logratios = (
            policy_preferred_logprob - policy_unpreferred_logprob
        )
        ref_logratios = (
            ref_preferred_logprob - ref_unpreferred_logprob
        )

        losses = -F.logsigmoid(
            self.beta * (policy_logratios - ref_logratios)
        )

        return losses.mean()


test.py

from __future__ import annotations
from typing import List, Tuple, Any, Optional

import torch
import torch.nn as nn
import torch.optim as optim
from torch import Tensor
from loguru import logger
from copy import deepcopy

# Define the end-of-sequence token ID (you should set this according to your tokenizer)
eos_token_id = 2  # Example EOS token ID

class Node:
    """
    A class representing a node in the tree of thoughts.

    Attributes:
        sequence (List[int]): The sequence of tokens from the root to this node.
        children (List[Node]): The list of child nodes.


import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
from loguru import logger
from typing import List, Tuple

# Set up logging
logger.add("training.log", rotation="500 MB")

# Device configuration
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class TransformerPolicyNetwork(nn.Module):
    """
    Transformer-based Policy Network that outputs action probabilities given a state sequence.
    """

    def __init__(
        self,


from collections.abc import Iterator
from copy import deepcopy

import torch
from torch.nn import Module
import torch.nn.functional as F
from x_transformers.x_transformers import TransformerWrapper

from einops import rearrange

# helper functions

def exists(v):
    return v is not None

def freeze_all_layers_(module):
    for param in module.parameters():
        param.requires_grad = False


from open_strawberry_torch.dpo import DPO

__all__ = ["DPO"]