Revert restore method to use original AbLang2 logic

adapter.py

@@ -561,48 +561,14 @@ 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]
         
-        restored_seqs = []
-        for seq in seqs:
-            # Tokenize the sequence using the correct interface
-            input_ids = self.tokenizer([seq], pad=True, w_extra_tkns=False, device=self.used_device)
-            
-            # Find masked tokens
-            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, return original sequence
-                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]  # [num_masked, vocab_size]
-                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 using the original tokenizer and remove spaces
-                restored_seq = self.tokenizer.tokenizer.decode(restored_input_ids, skip_special_tokens=True)
-                # Remove spaces that might be added during decoding
-                restored_seq = restored_seq.replace(' ', '')
-                restored_seqs.append(restored_seq)
-        
-        # For now, return restored sequences without alignment
-        # Proper ANARCI-based alignment requires full implementation of alignment functions
-        return np.array(restored_seqs) if len(restored_seqs) > 1 else restored_seqs[0]
+        # Use the original restore logic from the parent class
+        # The AbRestore class should have the working implementation
+        return super().restore(seqs, align=align, **kwargs)
 
 def add_angle_brackets(seq):
     # Assumes input is 'VH|VL' or 'VH|' or '|VL'