Fix restore method to properly handle masked tokens

adapter.py

@@ -561,14 +561,51 @@ class HFAbRestore(AbRestore):
     
     def restore(self, seqs, align=False, **kwargs):
         """Restore masked residues in antibody sequences."""
-        # Use the original AbLang2 restore logic
-        # This should work correctly like it did before
         if isinstance(seqs, str):
             seqs = [seqs]
         
-        # Use the original restore logic from the parent class
-        # The AbRestore class should have the working implementation
-        return super().restore(seqs, align=align, **kwargs)
+        restored_seqs = []
+        for seq in seqs:
+            # Check if sequence has masked tokens
+            if '*' not in seq:
+                # No masked tokens, return as-is
+                restored_seqs.append(seq)
+                continue
+            
+            # Tokenize the sequence
+            input_ids = self.tokenizer([seq], pad=True, w_extra_tkns=False, device=self.used_device)
+            
+            # 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]
+            
+            if len(masked_positions) == 0:
+                # No masked tokens found, return original
+                restored_seqs.append(seq)
+                continue
+            
+            # Get predictions for masked positions
+            with torch.no_grad():
+                output = self._hf_model(input_ids)
+                if hasattr(output, 'last_hidden_state'):
+                    logits = output.last_hidden_state
+                else:
+                    logits = output
+                
+                # Get predictions for masked positions
+                masked_logits = logits[0, masked_positions]
+                predicted_tokens = torch.argmax(masked_logits, dim=-1)
+                
+                # Replace masked tokens with predicted tokens
+                restored_input_ids = input_ids[0].clone()
+                restored_input_ids[masked_positions] = predicted_tokens
+                
+                # Decode back to sequence
+                restored_seq = self.tokenizer.tokenizer.decode(restored_input_ids, skip_special_tokens=True)
+                restored_seq = restored_seq.replace(' ', '')
+                restored_seqs.append(restored_seq)
+        
+        return np.array(restored_seqs) if len(restored_seqs) > 1 else restored_seqs[0]
 
 def add_angle_brackets(seq):
     # Assumes input is 'VH|VL' or 'VH|' or '|VL'