Add debug prints to restore method

adapter.py

@@ -572,15 +572,23 @@ class HFAbRestore(AbRestore):
                 restored_seqs.append(seq)
                 continue
             
+            # Debug: Print the sequence and mask token info
+            print(f"Processing sequence: {seq}")
+            print(f"Mask token ID: {self.tokenizer.mask_token_id}")
+            
             # Tokenize the sequence
             input_ids = self.tokenizer([seq], pad=True, w_extra_tkns=False, device=self.used_device)
+            print(f"Input IDs shape: {input_ids.shape}")
+            print(f"Input IDs: {input_ids[0]}")
             
             # Find masked tokens (assuming * is the mask token)
             mask_token_id = self.tokenizer.mask_token_id
             masked_positions = (input_ids[0] == mask_token_id).nonzero(as_tuple=True)[0]
+            print(f"Masked positions: {masked_positions}")
             
             if len(masked_positions) == 0:
                 # No masked tokens found, return original
+                print("No masked tokens found, returning original")
                 restored_seqs.append(seq)
                 continue
             
@@ -595,6 +603,7 @@ class HFAbRestore(AbRestore):
                 # Get predictions for masked positions
                 masked_logits = logits[0, masked_positions]
                 predicted_tokens = torch.argmax(masked_logits, dim=-1)
+                print(f"Predicted tokens: {predicted_tokens}")
                 
                 # Replace masked tokens with predicted tokens
                 restored_input_ids = input_ids[0].clone()
@@ -603,6 +612,7 @@ class HFAbRestore(AbRestore):
                 # Decode back to sequence
                 restored_seq = self.tokenizer.tokenizer.decode(restored_input_ids, skip_special_tokens=True)
                 restored_seq = restored_seq.replace(' ', '')
+                print(f"Restored sequence: {restored_seq}")
                 restored_seqs.append(restored_seq)
         
         return np.array(restored_seqs) if len(restored_seqs) > 1 else restored_seqs[0]