Update flow_matching/utils/multi_guidance.py

flow_matching/utils/multi_guidance.py

@@ -1,7 +1,8 @@
 import torch
-from flow_matching.utils import categorical
+import random
 import math
 import inspect
+import pdb
 
 def generate_simplex_lattice_points(num_obj: int, num_div: int) -> torch.Tensor:
     def rec(n, H):
@@ -28,13 +29,17 @@ def z_score_norm(tensor, eps=1e-8):
     std = tensor.std(dim=-1, unbiased=False, keepdim=True).clamp(min=eps)
     return (tensor - mean) / std
 
-def guided_transition_scoring(x_t, u_t, w, s_models, t, importance, args):
+def guided_transition_scoring(x_t, u_t, w, s_models, t, importance, tokenizer, args, fixed_positions=None, invalid_tokens=None):
     B, L, vocab_size = u_t.shape
     device = x_t.device
     guided_u_t = u_t.clone()  
     
     # 1. Randomly select one position per sequence.
-    pos_indices = torch.randint(low=1, high=L-2, size=(B,), device=device)  # shape: (B,)   # CHANGE!
+    all_positions = set(range(1, L-1))
+    available_positions = list(all_positions - set(fixed_positions))
+    assert len(available_positions) > 0
+    pos_indices = torch.tensor(random.choices(available_positions, k=B), device=device)
+    # pos_indices = torch.randint(low=1, high=L-2, size=(B,), device=device)  # shape: (B,)   # CHANGE!
     batch_idx = torch.arange(B, device=device)
     current_tokens = x_t[batch_idx, pos_indices]  # shape: (B,)
 
@@ -53,32 +58,42 @@ def guided_transition_scoring(x_t, u_t, w, s_models, t, importance, args):
     improvements_list = []
     with torch.no_grad():
         count = 0
+        input_seqs_cand = tokenizer.batch_decode(new_x_flat)
+        input_seqs_orig = tokenizer.batch_decode(x_t)
+        input_seqs_cand = [seq.replace(' ', '')[5:-5] for seq in input_seqs_cand]
+        input_seqs_orig = [seq.replace(' ', '')[5:-5] for seq in input_seqs_orig]
+        
         for i, s in enumerate(s_models):
             sig = inspect.signature(s.forward) if hasattr(s, 'forward') else inspect.signature(s)
             if 't' in sig.parameters:
-                candidate_scores = s(new_x_flat, t)
-                base_score = s(x_t, t)
+                candidate_scores = s(input_seqs_cand, t)
+                base_score = s(input_seqs_orig, t)
             else:
-                candidate_scores = s(new_x_flat)
-                base_score = s(x_t)
+                candidate_scores = s(input_seqs_cand)
+                base_score = s(input_seqs_orig)
 
             if isinstance(candidate_scores, tuple):
                 for k, score in enumerate(candidate_scores):
                     improvement = candidate_scores[k].view(B, vocab_size - 1) - base_score[k].unsqueeze(1)
-                    improvement = improvement.float()
+                    improvement = improvement.float().to(device)
                     improvement *= importance[count]
                     improvements_list.append(improvement.unsqueeze(2))
                     count += 1
             else:
                 improvement = candidate_scores.view(B, vocab_size - 1) - base_score.unsqueeze(1)
-                improvement = improvement.float()
+                improvement = improvement.float().to(device)
                 improvement *= importance[count]
                 improvements_list.append(improvement.unsqueeze(2))  # (B, vocab_size-1, 1)
                 count += 1
 
     improvement_values = torch.cat(improvements_list, dim=2) # (B, vocab_size-1, N)
-    if args.is_peptide:
-        improvement_values[:, :4, :] = -10 # Mask non-residue positions
+    
+    invalid_mask = cand_tokens.unsqueeze(-1) == invalid_tokens.view(1, 1, -1)
+    final_invalid_mask = invalid_mask.any(dim=-1)
+    improvement_values[final_invalid_mask] = -10.0
+
+    # if args.is_peptide:
+    #     improvement_values[:, :4, :] = -10 # Mask non-residue positions
 
     # 5. Compute ranking scores I_n
     ranks = torch.argsort(torch.argsort(improvement_values, dim=1), dim=1).float() + 1  # (B, vocab_size-1, N)
@@ -107,6 +122,98 @@ def guided_transition_scoring(x_t, u_t, w, s_models, t, importance, args):
 
     return guided_u_t, pos_indices, cand_tokens, improvement_values, delta_S
 
+
+def guided_transition_scoring_uaa(x_t, u_t, w, s_models, t, importance, tokenizer, args, fixed_positions=None, invalid_tokens=None):
+    B, L, vocab_size = u_t.shape
+    device = x_t.device
+    guided_u_t = u_t.clone()  
+    
+    # 1. Randomly select one position per sequence.
+    all_positions = set(range(1, L-1))
+    available_positions = list(all_positions - set(fixed_positions))
+    assert len(available_positions) > 0
+    pos_indices = torch.tensor(random.choices(available_positions, k=B), device=device)
+    # pos_indices = torch.randint(low=1, high=L-2, size=(B,), device=device)  # shape: (B,)   # CHANGE!
+    batch_idx = torch.arange(B, device=device)
+    current_tokens = x_t[batch_idx, pos_indices]  # shape: (B,)
+
+    # 2. Build candidate tokens for each sequence and remove self-transition.
+    full_cand_tokens = torch.arange(vocab_size, device=device).unsqueeze(0).expand(B, vocab_size) # (B, vocab_size)
+    mask = (full_cand_tokens != current_tokens.unsqueeze(1))  # (B, vocab_size)
+    # Now, cand_tokens contains only candidate tokens that differ from the current token.
+    cand_tokens = torch.masked_select(full_cand_tokens, mask).view(B, vocab_size - 1)  # (B, vocab_size-1)
+
+    # 3. Create candidate sequences by replacing the token at the selected position.
+    new_x = x_t.unsqueeze(1).expand(B, vocab_size, L).clone()
+    new_x = new_x[mask].view(B, vocab_size - 1, L)  # (B, vocab_size-1, L)
+    new_x[batch_idx, :, pos_indices] = cand_tokens 
+    new_x_flat = new_x.view(B * (vocab_size - 1), L)
+    improvements_list = []
+    with torch.no_grad():
+        count = 0
+        input_seqs_cand_smiles, valid_mask_cand = tokenizer.batch_decode(new_x_flat, convert_to_smiles=True, cyclic=args.cyclic)
+        input_seqs_cand_aa = tokenizer.batch_decode(new_x_flat, convert_to_smiles=False)
+
+        input_seqs_orig_smiles, valid_mask_orig = tokenizer.batch_decode(x_t, convert_to_smiles=True, cyclic=args.cyclic)
+        input_seqs_orig_aa = tokenizer.batch_decode(x_t, convert_to_smiles=False)
+        
+        for i, s in enumerate(s_models):
+            if i == 0:
+                candidate_scores = s(input_seqs_cand_aa) * valid_mask_cand
+                base_score = s(input_seqs_orig_aa) * valid_mask_orig
+            else:
+                candidate_scores = s(input_seqs_cand_smiles) * valid_mask_cand
+                base_score = s(input_seqs_orig_smiles) * valid_mask_orig
+
+            if isinstance(candidate_scores, tuple):
+                for k, score in enumerate(candidate_scores):
+                    improvement = candidate_scores[k].view(B, vocab_size - 1) - base_score[k].unsqueeze(1)
+                    improvement = improvement.float().to(device)
+                    improvement *= importance[count]
+                    improvements_list.append(improvement.unsqueeze(2))
+                    count += 1
+            else:
+                improvement = candidate_scores.view(B, vocab_size - 1) - base_score.unsqueeze(1)
+                improvement = improvement.float().to(device)
+                improvement *= importance[count]
+                improvements_list.append(improvement.unsqueeze(2))  # (B, vocab_size-1, 1)
+                count += 1
+
+    improvement_values = torch.cat(improvements_list, dim=2) # (B, vocab_size-1, N)
+    
+    invalid_mask = cand_tokens.unsqueeze(-1) == invalid_tokens.view(1, 1, -1)
+    final_invalid_mask = invalid_mask.any(dim=-1)
+    improvement_values[final_invalid_mask] = -10.0
+    
+
+    # 5. Compute ranking scores I_n
+    ranks = torch.argsort(torch.argsort(improvement_values, dim=1), dim=1).float() + 1  # (B, vocab_size-1, N)
+    I_n = ranks / float(vocab_size - 1)
+    avg_I = I_n.mean(dim=2)
+    norm_avg_I = z_score_norm(avg_I)    # (B, vocab_size-1)
+    
+    # 6. Compute directional score D
+    D = (improvement_values * w.view(1, 1, -1)).sum(dim=2)  
+    norm_D = z_score_norm(D)    # (B, vocab_size-1)
+
+    # 7. Combine the scores
+    delta_S = norm_avg_I + args.lambda_ * norm_D  # (B, vocab_size-1)
+
+    # 9. Update the guided velocities at the selected positions.
+    factor = torch.exp(args.beta * delta_S)  # (B, vocab_size-1)
+    factor = torch.clamp(factor, min=-100, max=100)
+
+    guided_u_t[batch_idx.unsqueeze(1), pos_indices.unsqueeze(1), cand_tokens] = u_t[batch_idx.unsqueeze(1), pos_indices.unsqueeze(1), cand_tokens] * factor
+
+    # 10. For the self-transition (current token) at the selected position, 
+    # set its guided velocity to be the negative sum of the updated off-diagonals.
+    updated_vals = guided_u_t[batch_idx, pos_indices, :]  # (B, vocab_size)
+    sum_off_diag = updated_vals.sum(dim=1) - updated_vals[batch_idx, current_tokens]
+    guided_u_t[batch_idx, pos_indices, current_tokens] = -sum_off_diag
+
+    return guided_u_t, pos_indices, cand_tokens, improvement_values, delta_S
+
+
 def adaptive_hypercone_filtering(improvement_values, cand_tokens, delta_S, w, Phi, args, ema_r_t=None):
     B, num_candidates, N = improvement_values.shape
     device = improvement_values.device