Add detailed debug prints to restore method

adapter.py

@@ -572,18 +572,31 @@ class HFAbRestore(AbRestore):
             print("WARNING: Alignment not implemented, skipping...")
             pass
         
+        print(f"DEBUG: Processing sequences: {seqs}")
+        
         # Tokenize sequences using original interface
         tokens = self.tokenizer(seqs, pad=True, w_extra_tkns=False, device=self.used_device)
+        print(f"DEBUG: Tokenized shape: {tokens.shape}")
+        print(f"DEBUG: First sequence tokens: {tokens[0]}")
         
         # Get predictions for amino acids (indices 1-20)
         predictions = self.AbLang(tokens)[:,:,1:21]
+        print(f"DEBUG: Predictions shape: {predictions.shape}")
         
         # Find predicted tokens and replace mask tokens
         predicted_tokens = torch.max(predictions, -1).indices + 1
+        print(f"DEBUG: Predicted tokens: {predicted_tokens[0]}")
+        
+        # Find mask token positions
+        mask_positions = (tokens == 23).nonzero(as_tuple=True)
+        print(f"DEBUG: Mask token positions: {mask_positions}")
+        
         restored_tokens = torch.where(tokens==23, predicted_tokens, tokens)
+        print(f"DEBUG: Restored tokens: {restored_tokens[0]}")
         
         # Decode back to sequences using original tokenizer
         restored_seqs = self.tokenizer(restored_tokens, mode="decode")
+        print(f"DEBUG: Decoded sequences: {restored_seqs}")
         
         # Handle paired sequences format
         if n_seqs < len(restored_seqs):
@@ -592,7 +605,9 @@ class HFAbRestore(AbRestore):
         
         # Apply final formatting
         from extra_utils import res_to_seq
-        return np.array([res_to_seq(seq, 'restore') for seq in np.c_[restored_seqs, np.vectorize(len)(seqs)]])
+        result = np.array([res_to_seq(seq, 'restore') for seq in np.c_[restored_seqs, np.vectorize(len)(seqs)]])
+        print(f"DEBUG: Final result: {result}")
+        return result
 
 def add_angle_brackets(seq):
     # Assumes input is 'VH|VL' or 'VH|' or '|VL'
@@ -939,16 +954,13 @@ class AbLang2PairedHuggingFaceAdapter(AbEncoding, AbRestore, AbAlignment, AbScor
             return [res_to_list(state, seq) for state, seq in zip(probs, formatted_seqs)]
 
     def restore(self, seqs, align=False, **kwargs):
-        print(f"DEBUG: Input sequences: {seqs}")
         hf_abrestore = HFAbRestore(self.AbLang, self.tokenizer, spread=self.spread, device=self.used_device, ncpu=self.ncpu)
         restored = hf_abrestore.restore(seqs, align=align)
-        print(f"DEBUG: Restored before formatting: {restored}")
         # Apply angle brackets formatting
         if isinstance(restored, np.ndarray):
             restored = np.array([add_angle_brackets(seq) for seq in restored])
         else:
             restored = [add_angle_brackets(seq) for seq in restored]
-        print(f"DEBUG: Final output: {restored}")
         return restored
 
 def extract_input_ids(tokens, device):