Upload 6 files

utils/dataloader.py

@@ -0,0 +1,41 @@
+import torch
+from functools import partial
+from torch.utils.data import DataLoader
+from torch import nn
+
+def collate_fn(batch):
+  input_ids = torch.tensor(batch[0]['input_ids'])
+  attention_mask = torch.tensor(batch[0]['attention_mask'])
+  return {
+    'input_ids': input_ids,
+    'attention_mask': attention_mask
+  }
+
+class CustomDataModule(nn.Module):
+    def __init__(self, train_dataset, val_dataset, test_dataset, collate_fn=collate_fn):
+        super().__init__()
+        self.train_dataset = train_dataset
+        self.val_dataset = val_dataset
+        self.test_dataset = test_dataset
+        self.collate_fn = collate_fn
+
+    def train_dataloader(self):
+        return DataLoader(self.train_dataset, 
+                          collate_fn=partial(self.collate_fn),
+                          num_workers=8,
+                          pin_memory=True,
+                          shuffle=True)
+    
+    def val_dataloader(self):
+        return DataLoader(self.val_dataset, 
+                          collate_fn=partial(self.collate_fn),
+                          num_workers=8,
+                          pin_memory=True,
+                          shuffle=False)
+  
+    def test_dataloader(self):
+        return DataLoader(self.test_dataset, 
+                          collate_fn=partial(self.collate_fn),
+                          num_workers=8,
+                          pin_memory=True,
+                          shuffle=False)

utils/dataset.py

@@ -0,0 +1,17 @@
+import copy
+import pickle
+import torch
+
+class EnhancerDataset(torch.utils.data.Dataset):
+    def __init__(self, mel_enhancer=True, split='train'):
+        all_data = pickle.load(open(f'./dataset/enhancer_data/Deep{"MEL2" if mel_enhancer else "FlyBrain"}_data.pkl', 'rb'))
+        self.seqs = torch.argmax(torch.from_numpy(copy.deepcopy(all_data[f'{split}_data'])), dim=-1)
+        self.clss = torch.argmax(torch.from_numpy(copy.deepcopy(all_data[f'y_{split}'])), dim=-1)
+        self.num_cls = all_data[f'y_{split}'].shape[-1]
+        self.alphabet_size = 4
+
+    def __len__(self):
+        return len(self.seqs)
+
+    def __getitem__(self, idx):
+        return self.seqs[idx], self.clss[idx]

utils/flow_utils.py

@@ -0,0 +1,497 @@
+import copy
+import math
+import pickle
+
+import scipy
+import torch.nn.functional as F
+import numpy as np
+import torch
+import torch.nn as nn
+from scipy.linalg import sqrtm
+import re
+
+
+
+
+
+def upgrade_state_dict(state_dict, prefixes=["encoder.sentence_encoder.", "encoder."]):
+    """Removes prefixes 'model.encoder.sentence_encoder.' and 'model.encoder.'."""
+    pattern = re.compile("^" + "|".join(prefixes))
+    state_dict = {pattern.sub("", name): param for name, param in state_dict.items()}
+    return state_dict
+
+def map_t_to_alpha(t, alpha_scale):
+    """
+    Maps t in [0,1) to the range of alphas using the inverse CDF of an exponential distribution.
+
+    Args:
+        t (torch.Tensor): A tensor of values in [0,1).
+        alpha_scale (float): The scaling factor used in the original alpha calculation.
+
+    Returns:
+        torch.Tensor: The corresponding alpha values.
+    """
+    if torch.any(t >= 1) or torch.any(t < 0):
+        raise ValueError("t must be in the range [0,1).")
+    
+    return 1 + (-torch.log(1 - t)) * alpha_scale
+    
+    # return torch.clamp(1 + (-torch.log(1 - t)) * alpha_scale, torch.tensor(8).to(t.device))
+
+def load_flybrain_designed_seqs(path):
+    order = {'A': 0, 'C':1, 'G':2, 'T':3}
+    f = open(path, "rb")
+    data = pickle.load(f)
+    arrays = []
+    for seq in data['seq']:
+        arrays.append([order[char] for char in seq])
+    return torch.tensor(arrays, dtype=torch.long)
+
+
+def update_ema(current_dict, prev_ema, gamma = 0.9):
+    ema = copy.deepcopy(prev_ema)
+    current_dict = copy.deepcopy(current_dict)
+    for key, current_value in current_dict.items():
+        ema_key  = 'ema_' + key
+        if not np.isnan(current_value):
+            if ema_key in prev_ema:
+                ema[ema_key] = (1 - gamma) * current_value + gamma * prev_ema[ema_key]
+            else:
+                ema[ema_key] = current_value
+    return ema
+
+def min_max_str(x):
+    return f'min {x.min()} max {x.max()}'
+
+def get_wasserstein_dist(embeds1, embeds2):
+    if np.isnan(embeds2).any() or np.isnan(embeds1).any() or len(embeds1) == 0 or len(embeds2) == 0:
+        return float('nan')
+    mu1, sigma1 = embeds1.mean(axis=0), np.cov(embeds1, rowvar=False)
+    mu2, sigma2 = embeds2.mean(axis=0), np.cov(embeds2, rowvar=False)
+    ssdiff = np.sum((mu1 - mu2) ** 2.0)
+    covmean = sqrtm(sigma1.dot(sigma2))
+    if np.iscomplexobj(covmean):
+        covmean = covmean.real
+    dist = ssdiff + np.trace(sigma1 + sigma2 - 2.0 * covmean)
+    return dist
+
+def simplex_proj(seq):
+    """Algorithm from https://arxiv.org/abs/1309.1541 Weiran Wang, Miguel Á. Carreira-Perpiñán"""
+    Y = seq.reshape(-1, seq.shape[-1])
+    N, K = Y.shape
+    X, _ = torch.sort(Y, dim=-1, descending=True)
+    X_cumsum = torch.cumsum(X, dim=-1) - 1
+    div_seq = torch.arange(1, K + 1, dtype=Y.dtype, device=Y.device)
+    Xtmp = X_cumsum / div_seq.unsqueeze(0)
+
+    greater_than_Xtmp = (X > Xtmp).sum(dim=1, keepdim=True)
+    row_indices = torch.arange(N, dtype=torch.long, device=Y.device).unsqueeze(1)
+    selected_Xtmp = Xtmp[row_indices, greater_than_Xtmp - 1]
+
+    X = torch.max(Y - selected_Xtmp, torch.zeros_like(Y))
+    return X.view(seq.shape)
+
+
+
+def batch_project_simplex(v):
+    u, _ = torch.sort(v, dim=1, descending=True)
+    cssv = u.cumsum(dim=1)
+    k = torch.arange(1, v.shape[1] + 1, device=v.device)
+    rho = ((u * k) > (cssv - 1)).int().cumsum(dim=1).argmax(dim=1)
+    theta = (cssv[torch.arange(v.shape[0]), rho] - 1) / (rho + 1).float()
+    w = torch.maximum(v - theta.unsqueeze(1), torch.tensor(0.0, device=v.device))
+    return w
+
+if __name__ == "__main__":
+    a = torch.softmax(torch.rand((5,4)), dim=-1)
+    b = torch.rand((5,4)) - 1
+    ab = torch.cat([a,b])
+    ab_proj1 = batch_project_simplex(ab)
+    ab_proj2 = simplex_proj(ab)
+    print('ab_proj1 - ab_proj2',ab_proj1 - ab_proj2)
+    print('ab_proj1 - ab', ab_proj1 - ab)
+    print('ab_proj2.sum(-1)', ab_proj2.sum(-1))
+    print('ab_proj2', ab_proj2)
+
+def sample_cond_prob_path(args, seq, alphabet_size):
+    B, L = seq.shape
+    seq_one_hot = torch.nn.functional.one_hot(seq, num_classes=alphabet_size)
+    if args.mode == 'dirichlet':
+        alphas = torch.from_numpy(1 + scipy.stats.expon().rvs(size=B) * args.alpha_scale).to(seq.device).float()
+        if args.fix_alpha:
+            alphas = torch.ones(B, device=seq.device) * args.fix_alpha
+        alphas_ = torch.ones(B, L, alphabet_size, device=seq.device)
+        alphas_ = alphas_ + seq_one_hot * (alphas[:,None,None] - 1)
+        xt = torch.distributions.Dirichlet(alphas_).sample()
+    elif args.mode == 'distill':
+        alphas = torch.zeros(B, device=seq.device)
+        xt = torch.distributions.Dirichlet(torch.ones(B, L, alphabet_size, device=seq.device)).sample()
+    elif args.mode == 'riemannian':
+        t = torch.rand(B, device=seq.device)
+        dirichlet = torch.distributions.Dirichlet(torch.ones(alphabet_size, device=seq.device))
+        x0 = dirichlet.sample((B,L))
+        x1 = seq_one_hot
+        xt = t[:,None,None] * x1 + (1 - t[:,None,None]) * x0
+        alphas = t
+    elif args.mode == 'ardm' or args.mode == 'lrar':
+        mask_prob = torch.rand(1, device=seq.device)
+        mask = torch.rand(seq.shape, device=seq.device) < mask_prob
+        if args.mode == 'lrar': mask = ~(torch.arange(L, device=seq.device) < (1-mask_prob) * L)
+        xt = torch.where(mask, alphabet_size, seq) # mask token index
+        xt = torch.nn.functional.one_hot(xt, num_classes=alphabet_size + 1).float() # plus one to include index for mask token
+        alphas = mask_prob.expand(B)
+    return xt, alphas
+
+def expand_simplex(xt, alphas, prior_pseudocount):
+    prior_weights = (prior_pseudocount / (alphas + prior_pseudocount - 1))[:, None, None]
+    return torch.cat([xt * (1 - prior_weights), xt * prior_weights], -1), prior_weights
+
+
+class DirichletConditionalFlow:
+    def __init__(self, K=20, alpha_min=1, alpha_max=100, alpha_spacing=0.01):
+        self.alphas = np.arange(alpha_min, alpha_max + alpha_spacing, alpha_spacing)
+        self.beta_cdfs = []
+        self.bs = np.linspace(0, 1, 1000)
+        for alph in self.alphas:
+            self.beta_cdfs.append(scipy.special.betainc(alph, K-1, self.bs))
+        self.beta_cdfs = np.array(self.beta_cdfs)
+        self.beta_cdfs_derivative = np.diff(self.beta_cdfs, axis=0) / alpha_spacing
+        self.K = K
+
+    def c_factor(self, bs, alpha):
+        out1 = scipy.special.beta(alpha, self.K - 1)
+        out2 = np.where(bs < 1, out1 / ((1 - bs) ** (self.K - 1)), 0)
+        out = np.where((bs ** (alpha - 1)) > 0, out2 / (bs ** (alpha - 1)), 0)
+        I_func = self.beta_cdfs_derivative[np.argmin(np.abs(alpha - self.alphas))]
+        interp = -np.interp(bs, self.bs, I_func)
+        final = interp * out
+        return final
+
+
+class GaussianSmearing(torch.nn.Module):
+    # used to embed the edge distances
+    def __init__(self, start=0.0, stop=5.0, embedding_dim=50):
+        super().__init__()
+        offset = torch.linspace(start, stop, embedding_dim)
+        self.coeff = -0.5 / (offset[1] - offset[0]).item() ** 2
+        self.register_buffer("offset", offset)
+        self.embedding_dim = embedding_dim
+
+    def forward(self, signal):
+        shape = signal.shape
+        signal = signal.view(-1, 1) - self.offset.view(1, -1) + 1E-6
+        encoded = torch.exp(self.coeff * torch.pow(signal, 2))
+        return encoded.view(*shape, self.embedding_dim)
+
+
+class MonotonicFunction(torch.nn.Module):
+    def __init__(self, init_max, num_bins):
+        super().__init__()
+        self.w = torch.nn.Parameter(torch.ones(num_bins) * np.log(init_max) - np.log(num_bins))
+        self.num_bins = num_bins
+
+    def forward(self, t):
+        widths = torch.exp(self.w)
+        right = torch.cumsum(widths, 0)
+        left = right - widths
+
+        bin_idx = (t * self.num_bins).long()
+        frac_part = t - bin_idx * (1 / self.num_bins)
+
+        return left[bin_idx] + (frac_part * self.num_bins) * (right[bin_idx] - left[bin_idx])
+
+    def invert(self, f):
+        widths = torch.exp(self.w)
+        left = torch.cumsum(widths, 0) - widths
+        bin_idx = (f.unsqueeze(-1) > left).sum(-1) - 1
+        frac_part = f - left[bin_idx]
+        return bin_idx / self.num_bins + frac_part / widths[bin_idx] / self.num_bins
+
+    def derivative(self, t):
+        widths = torch.exp(self.w)
+        right = torch.cumsum(widths, 0)
+        left = right - widths
+        bin_idx = (t * self.num_bins).long()
+        return (right[bin_idx] - left[bin_idx]) * self.num_bins
+
+class SinusoidalEmbedding(nn.Module):
+    """ from https://github.com/hojonathanho/diffusion/blob/master/diffusion_tf/nn.py   """
+    def __init__(self, embedding_dim, embedding_scale, max_positions=10000):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.max_positions = max_positions
+        self.embedding_scale = embedding_scale
+
+    def forward(self, signal):
+        shape = signal.shape
+        signal = signal.view(-1) * self.embedding_scale
+        half_dim = self.embedding_dim // 2
+        emb = math.log(self.max_positions) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, dtype=torch.float32, device=signal.device) * -emb)
+        emb = signal.float()[:, None] * emb[None, :]
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+        if self.embedding_dim % 2 == 1:  # zero pad
+            emb = F.pad(emb, (0, 1), mode='constant')
+        assert emb.shape == (signal.shape[0], self.embedding_dim)
+        return emb.view(*shape, self.embedding_dim )
+
+
+class GaussianFourierProjection(nn.Module):
+    """Gaussian Fourier embeddings for noise levels.
+    from https://github.com/yang-song/score_sde_pytorch/blob/1618ddea340f3e4a2ed7852a0694a809775cf8d0/models/layerspp.py#L32
+    """
+
+    def __init__(self, embedding_dim=256, scale=1.0):
+        super().__init__()
+        self.W = nn.Parameter(torch.randn(embedding_dim//2) * scale, requires_grad=False)
+        self.embedding_dim = embedding_dim
+
+    def forward(self, signal):
+        shape = signal.shape
+        signal = signal.view(-1)
+        signal_proj = signal[:, None] * self.W[None, :] * 2 * np.pi
+        emb = torch.cat([torch.sin(signal_proj), torch.cos(signal_proj)], dim=-1)
+        return emb.view(*shape, self.embedding_dim )
+
+def get_signal_mapping(embedding_type, embedding_dim, embedding_scale=10000):
+    if embedding_type == 'sinusoidal':
+        emb_func = SinusoidalEmbedding(embedding_dim=embedding_dim, embedding_scale=embedding_scale)
+    elif embedding_type == 'fourier':
+        emb_func = GaussianFourierProjection(embedding_dim=embedding_dim, scale=embedding_scale)
+    elif embedding_type == 'gaussian':
+        emb_func = GaussianSmearing(0.0, 1, embedding_dim)
+    else:
+        raise NotImplemented
+    return emb_func
+
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function,
+    which defines the cumulative product of (1-beta) over time from t = [0,1].
+
+    :param num_diffusion_timesteps: the number of betas to produce.
+    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
+                      produces the cumulative product of (1-beta) up to that
+                      part of the diffusion process.
+    :param max_beta: the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+    """
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    return np.array(betas)
+
+def get_beta_schedule(num_steps):
+
+    return betas_for_alpha_bar(
+            num_steps,
+            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
+        )
+
+
+class GaussianDiffusionSchedule:
+    """
+    Utilities for training and sampling diffusion models.
+
+    Ported directly from here, and then adapted over time to further experimentation.
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42
+
+    :param betas: a 1-D numpy array of betas for each diffusion timestep,
+                  starting at T and going to 1.
+    :param model_mean_type: a ModelMeanType determining what the model outputs.
+    :param model_var_type: a ModelVarType determining how variance is output.
+    :param loss_type: a LossType determining the loss function to use.
+    :param rescale_timesteps: if True, pass floating point timesteps into the
+                              model so that they are always scaled like in the
+                              original paper (0 to 1000).
+    """
+
+    def __init__(
+            self,
+            timesteps,
+            noise_scale=1.0,
+    ):
+        betas = get_beta_schedule(timesteps)
+
+        # Use float64 for accuracy.
+        betas = np.array(betas, dtype=np.float64)
+        self.betas = betas
+        assert len(betas.shape) == 1, "betas must be 1-D"
+        assert (betas > 0).all() and (betas <= 1).all()
+
+        self.timesteps = int(betas.shape[0])
+        self.noise_scale = noise_scale
+
+        alphas = 1.0 - betas
+        self.alphas_cumprod = np.cumprod(alphas, axis=0)
+        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
+        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
+        assert self.alphas_cumprod_prev.shape == (self.timesteps,)
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
+        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
+        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
+        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
+        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        self.posterior_variance = (
+                betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        # log calculation clipped because the posterior variance is 0 at the
+        # beginning of the diffusion chain.
+        self.posterior_log_variance_clipped = np.log(
+            np.append(self.posterior_variance[1], self.posterior_variance[1:])
+        )
+        self.posterior_mean_coef1 = (
+                betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        self.posterior_mean_coef2 = (
+                (1.0 - self.alphas_cumprod_prev)
+                * np.sqrt(alphas)
+                / (1.0 - self.alphas_cumprod)
+        )
+
+    def q_sample(self, x_start, t, noise=None):
+        """
+        Diffuse the data for a given number of diffusion steps.
+
+        In other words, sample from q(x_t | x_0).
+
+        :param x_start: the initial data batch.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :param noise: if specified, the split-out normal noise.
+        :return: A noisy version of x_start.
+        """
+        if noise is None:
+            noise = self.noise_scale * torch.randn_like(x_start)
+            # add scaling here
+        assert noise.shape == x_start.shape
+        return (
+                _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+                + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
+                * noise
+        )
+
+    def q_posterior_mean_variance(self, x_start, x_t, t):
+        """
+        Compute the mean and variance of the diffusion posterior:
+
+            q(x_{t-1} | x_t, x_0)
+
+        """
+        assert x_start.shape == x_t.shape
+        posterior_mean = (
+                _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
+                + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = _extract_into_tensor(
+            self.posterior_log_variance_clipped, t, x_t.shape
+        )
+
+        posterior_variance = (self.noise_scale ** 2) * posterior_variance
+        posterior_log_variance_clipped = 2 * np.log(self.noise_scale) + posterior_log_variance_clipped
+
+        assert (
+                posterior_mean.shape[0]
+                == posterior_variance.shape[0]
+                == posterior_log_variance_clipped.shape[0]
+                == x_start.shape[0]
+        )
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+
+def _extract_into_tensor(arr, timesteps, broadcast_shape):
+    """
+    Extract values from a 1-D numpy array for a batch of indices.
+
+    :param arr: the 1-D numpy array.
+    :param timesteps: a tensor of indices into the array to extract.
+    :param broadcast_shape: a larger shape of K dimensions with the batch
+                            dimension equal to the length of timesteps.
+    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
+    """
+    res = torch.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
+    while len(res.shape) < len(broadcast_shape):
+        res = res[..., None]
+    return res.expand(broadcast_shape)
+
+
+def space_timesteps(num_timesteps, section_counts):
+    """
+    Create a list of timesteps to use from an original diffusion process,
+    given the number of timesteps we want to take from equally-sized portions
+    of the original process.
+
+    For example, if there's 300 timesteps and the section counts are [10,15,20]
+    then the first 100 timesteps are strided to be 10 timesteps, the second 100
+    are strided to be 15 timesteps, and the final 100 are strided to be 20.
+
+    If the stride is a string starting with "ddim", then the fixed striding
+    from the DDIM paper is used, and only one section is allowed.
+
+    :param num_timesteps: the number of diffusion steps in the original
+                          process to divide up.
+    :param section_counts: either a list of numbers, or a string containing
+                           comma-separated numbers, indicating the step count
+                           per section. As a special case, use "ddimN" where N
+                           is a number of steps to use the striding from the
+                           DDIM paper.
+    :return: a set of diffusion steps from the original process to use.
+    """
+    if isinstance(section_counts, str):
+        if section_counts.startswith("ddim"):
+            desired_count = int(section_counts[len("ddim"):])
+            for i in range(1, num_timesteps):
+                if len(range(0, num_timesteps, i)) == desired_count:
+                    return set(range(0, num_timesteps, i))
+            raise ValueError(
+                f"cannot create exactly {num_timesteps} steps with an integer stride"
+            )
+        section_counts = [int(x) for x in section_counts.split(",")]
+    size_per = num_timesteps // len(section_counts)
+    extra = num_timesteps % len(section_counts)
+    start_idx = 0
+    all_steps = []
+    for i, section_count in enumerate(section_counts):
+        size = size_per + (1 if i < extra else 0)
+        if size < section_count:
+            raise ValueError(
+                f"cannot divide section of {size} steps into {section_count}"
+            )
+        if section_count <= 1:
+            frac_stride = 1
+        else:
+            frac_stride = (size - 1) / (section_count - 1)
+        cur_idx = 0.0
+        taken_steps = []
+        for _ in range(section_count):
+            taken_steps.append(start_idx + round(cur_idx))
+            cur_idx += frac_stride
+        all_steps += taken_steps
+        start_idx += size
+    return set(all_steps)
+
+def timestep_embedding(timesteps, dim, max_period=10000):
+    """
+    Create sinusoidal timestep embeddings.
+
+    :param timesteps: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an [N x dim] Tensor of positional embeddings.
+    """
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+    ).to(device=timesteps.device)
+    args = timesteps[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    return embedding

utils/parsing.py

@@ -0,0 +1,30 @@
+from argparse import ArgumentParser
+import math
+
+def parse_guidance_args():
+    parser = ArgumentParser()
+    
+    parser.add_argument("--num_div", type=int, default=64)
+    parser.add_argument("--lambda_", type=float, default=1.0)
+    parser.add_argument("--beta", type=float, default=1.0)
+    parser.add_argument("--alpha_r", type=float, default=0.5)
+    parser.add_argument("--eta", type=float, default=1.0)
+    parser.add_argument("--Phi_init", type=float, default=math.radians(45.0))
+    parser.add_argument("--Phi_min", type=float, default=math.radians(15.0))
+    parser.add_argument("--Phi_max", type=float, default=math.radians(75.0))
+    parser.add_argument("--tau", type=float, default=0.3)
+    parser.add_argument("--T", type=int, default=100)
+    parser.add_argument("--length", type=int, default=12)
+    parser.add_argument("--is_peptide", type=bool, default=True)
+    parser.add_argument("--n_samples", type=int, default=5)
+    parser.add_argument("--n_batches", type=int, default=2)
+    parser.add_argument("--target_protein", type=str, default="AAAAA")
+    parser.add_argument("--target_enhancer_class", type=int, default=0)
+    parser.add_argument("--target_DNA_shape", type=str, default='HelT')
+    parser.add_argument("--motifs", type=str, required=False)
+    parser.add_argument("--weights", type=float, nargs='+', required=False)
+    parser.add_argument("--output_file", type=str, default='moo_outputs.txt')
+    parser.add_argument("--motif_penalty", action='store_true')
+    
+    args = parser.parse_args()
+    return args