Update flow_matching/solver/discrete_solver.py

flow_matching/solver/discrete_solver.py

@@ -10,7 +10,6 @@ from typing import Callable, Optional, Union
 
 import torch
 from torch import Tensor
-import gc
 from torch.nn import functional as F
 
 from flow_matching.path import MixtureDiscreteProbPath
@@ -21,7 +20,7 @@ from .utils import get_nearest_times
 from ..utils.multi_guidance import *
 
 try:
-    from tqdm import tqdm
+    from tqdm.auto import tqdm
 
     TQDM_AVAILABLE = True
 except ImportError:
@@ -275,18 +274,12 @@ class MixtureDiscreteEulerSolver(Solver):
         score_models: list = None,
         num_objectives: int = 1,
         weights: list = None,
+        tokenizer = None,
+        fixed_positions=None,
+        invalid_tokens=None,
         **model_extras,
     ) -> Tensor:
 
-        # score_list_0 = []
-        # score_list_1 = []
-        # score_list_2 = []
-        # score_list_3 = []
-        # score_list_4 = []
-        # score_list_5 = []
-
-        import pdb
-
         if not div_free == 0.0:
             raise NotImplementedError
 
@@ -331,7 +324,7 @@ class MixtureDiscreteEulerSolver(Solver):
                 raise ImportError(
                     "tqdm is required for verbose mode. Please install it."
                 )
-            ctx = tqdm(total=t_final, desc=f"NFE: {steps_counter}")
+            ctx = tqdm(total=n_steps, desc=f"NFE", dynamic_ncols=True, leave=True, bar_format="{desc}: {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt}{postfix}")
         else:
             ctx = nullcontext()
 
@@ -342,7 +335,7 @@ class MixtureDiscreteEulerSolver(Solver):
             w, _ = select_random_weight_vector(num_objectives, args.num_div)
             # w = torch.tensor([0.2, 0.7, 0.05, 0.05]).to(x_t.device)
             w = w.to(device=x_init.device)
-        print(f"Weight Vector: {w}")
+        # print(f"Weight Vector: {w}")
         Phi = args.Phi_init
         ema_r_t = None
 
@@ -362,14 +355,10 @@ class MixtureDiscreteEulerSolver(Solver):
                     d_k_t = scheduler_output.d_alpha_t
                     u_t = d_k_t / (1 - k_t) * p_1t
 
-                    guided_u_t, pos_indices, cand_tokens, improvement_values, delta_S = guided_transition_scoring(x_t, u_t, w, score_models, t, w, args)
+                    guided_u_t, pos_indices, cand_tokens, improvement_values, delta_S = guided_transition_scoring(x_t, u_t, w, score_models, t, w, tokenizer, args, fixed_positions, invalid_tokens)
 
                     best_candidate, accepted_mask, valid_mask, Phi, ema_r_t = adaptive_hypercone_filtering(improvement_values, cand_tokens, delta_S, w, Phi, args, ema_r_t=ema_r_t)
                     
-                    # best_candidate, accepted_mask, valid_mask, Phi, ema_r_t = hypercone_filtering(improvement_values, cand_tokens, delta_S, w, Phi, args, ema_r_t=ema_r_t)
-
-                    # best_candidate = get_best_candidate(improvement_values, cand_tokens, delta_S)
-
                     x_t = euler_sample(x_t, pos_indices, best_candidate, guided_u_t, h)
 
 
@@ -377,37 +366,165 @@ class MixtureDiscreteEulerSolver(Solver):
                 t = t + h
 
                 scores = []
+                input_seqs = tokenizer.batch_decode(x_t)
+                input_seqs = [seq.replace(' ', '')[5:-5] for seq in input_seqs]
+                
                 for i, s in enumerate(score_models):
                     sig = inspect.signature(s.forward) if hasattr(s, 'forward') else inspect.signature(s)
                     if 't' in sig.parameters:
-                        candidate_scores = s(x_t, 1)
+                        candidate_scores = s(input_seqs, 1)
                     else:
-                        candidate_scores = s(x_t)
+                        candidate_scores = s(input_seqs)
 
                     if isinstance(candidate_scores, tuple):
                         for score in candidate_scores:
                             scores.append(score.item())
                     else:
                         scores.append(candidate_scores.item())
-                print(scores)
-
-                    # print(f"Score {i}: {[round(s.item(), 4) for s in candidate_scores]}")
-                    # if i == 0:
-                    #     score_list_0.append(round(candidate_scores[0].item(), 2))
-                    #     # score_list_0.append(round(1-candidate_scores.item(), 2))
-                    #     # score_list_1.append(round(candidate_scores[1].item(), 2))
-                    # if i == 1:
-                    #     score_list_1.append(round(candidate_scores.item(), 2))
-                    #     # score_list_2.append(round(candidate_scores.item(), 2))
-                    # if i == 2:
-                    #     score_list_2.append(round(candidate_scores.item(), 2))
-                    # if i == 3:
-                    #     score_list_3.append(round(candidate_scores.item(), 2))
-                    # if i == 4:
-                    #     score_list_4.append(round(candidate_scores.item(), 2))
-                    # if i == 5:
-                    #     score_list_5.append(round(candidate_scores.item(), 2))
 
+                postfix = {}
+                for i, objective in enumerate(args.objectives):
+                    postfix[objective] = scores[i]
+                
+                ctx.set_description(f"NFE: {steps_counter}", refresh=False)
+                ctx.set_postfix({k: f"{v:.3f}" for k, v in postfix.items()}, refresh=False)
+                ctx.update(1)
+
+                if return_intermediates and (t in time_grid):
+                    res.append(x_t.clone())
+
+        if return_intermediates:
+            if step_size is None:
+                return torch.stack(res, dim=0)
+            else:
+                return torch.stack(res, dim=0)[order]
+        else:
+            return x_t
+
+
+    def multi_guidance_sample_uaa(
+        self,
+        args,
+        x_init: Tensor,
+        step_size: Optional[float],
+        div_free: Union[float, Callable[[float], float]] = 0.0,
+        dtype_categorical: torch.dtype = torch.float32,
+        time_grid: Tensor = torch.tensor([0.0, 1.0]),
+        return_intermediates: bool = False,
+        verbose: bool = False,
+        score_models: list = None,
+        num_objectives: int = 1,
+        weights: list = None,
+        tokenizer = None,
+        fixed_positions=None,
+        invalid_tokens=None,
+        **model_extras,
+    ) -> Tensor:
+
+        if not div_free == 0.0:
+            raise NotImplementedError
+
+        # Initialize the current state `x_t` with the initial state `X_0`.
+        time_grid = time_grid.to(device=x_init.device)
+
+        if step_size is None:
+            # If step_size is None then set the t discretization to time_grid.
+            t_discretization = time_grid
+            n_steps = len(time_grid) - 1
+        else:
+            # If step_size is float then t discretization is uniform with step size set by step_size.
+            t_init = time_grid[0].item()
+            t_final = time_grid[-1].item()
+            assert (
+                t_final - t_init
+            ) > step_size, f"Time interval [time_grid[0], time_grid[-1]] must be larger than step_size. Got a time interval [{t_init}, {t_final}] and step_size {step_size}."
+
+            n_steps = ceil((t_final - t_init) / step_size)
+            t_discretization = torch.tensor(
+                [t_init + step_size * i for i in range(n_steps)] + [t_final],
+                device=x_init.device,
+            )
+
+            if return_intermediates:
+                # get order of intermediate steps:
+                order = torch.argsort(time_grid)
+                # Compute intermediate steps to return via nearest points in t_discretization to time_grid.
+                time_grid = get_nearest_times(
+                    time_grid=time_grid, t_discretization=t_discretization
+                )
+
+        x_t = x_init.clone()
+        steps_counter = 0
+        res = []
+
+        if return_intermediates:
+            res = [x_init.clone()]
+
+        if verbose:
+            if not TQDM_AVAILABLE:
+                raise ImportError(
+                    "tqdm is required for verbose mode. Please install it."
+                )
+            ctx = tqdm(total=t_final, desc=f"NFE: {steps_counter}")
+        else:
+            ctx = nullcontext()
+
+        # Randomly sample a weight vector
+        if weights is not None:
+            w = torch.tensor(weights).to(device=x_init.device)
+        else:
+            w, _ = select_random_weight_vector(num_objectives, args.num_div)
+            # w = torch.tensor([0.2, 0.7, 0.05, 0.05]).to(x_t.device)
+            w = w.to(device=x_init.device)
+        # print(f"Weight Vector: {w}")
+        Phi = args.Phi_init
+        ema_r_t = None
+
+        with ctx:
+            for i in range(n_steps):
+                t = t_discretization[i : i + 1]
+                h = t_discretization[i + 1 : i + 2] - t_discretization[i : i + 1]
+
+                p_1t = self.model(x=x_t, t=t.repeat(x_t.shape[0]), **model_extras)
+                x_1 = categorical(p_1t.to(dtype=dtype_categorical))
+
+                # Checks if final step
+                if i != n_steps - 1:
+                    # Compute u_t(y,x)
+                    scheduler_output = self.path.scheduler(t=t)
+                    k_t = scheduler_output.alpha_t
+                    d_k_t = scheduler_output.d_alpha_t
+                    u_t = d_k_t / (1 - k_t) * p_1t
+
+                    guided_u_t, pos_indices, cand_tokens, improvement_values, delta_S = guided_transition_scoring_uaa(x_t, u_t, w, score_models, t, w, tokenizer, args, fixed_positions, invalid_tokens)
+
+                    best_candidate, accepted_mask, valid_mask, Phi, ema_r_t = adaptive_hypercone_filtering(improvement_values, cand_tokens, delta_S, w, Phi, args, ema_r_t=ema_r_t)
+                    
+                    # best_candidate, accepted_mask, valid_mask, Phi, ema_r_t = hypercone_filtering(improvement_values, cand_tokens, delta_S, w, Phi, args, ema_r_t=ema_r_t)
+
+                    # best_candidate = get_best_candidate(improvement_values, cand_tokens, delta_S)
+
+                    x_t = euler_sample(x_t, pos_indices, best_candidate, guided_u_t, h)
+
+                steps_counter += 1
+                t = t + h
+
+                scores = []
+                input_seqs_smiles, _ = tokenizer.batch_decode(x_t, convert_to_smiles=True, cyclic=args.cyclic)
+                input_seqs_aa = tokenizer.batch_decode(x_t, convert_to_smiles=False)
+                
+                for i, s in enumerate(score_models):
+                    if i == 0:
+                        score = s(input_seqs_aa)
+                    else:
+                        score = s(input_seqs_smiles)
+                    
+                    if isinstance(score, tuple):
+                        for s in score:
+                            scores.append(s.item())
+                    else:
+                        scores.append(score.item())
+                ctx.write(scores)
 
                 if return_intermediates and (t in time_grid):
                     res.append(x_t.clone())