Implement working restore method in HFAbRestore class

adapter.py

@@ -559,6 +559,40 @@ class HFAbRestore(AbRestore):
                 return output.last_hidden_state
             return output
         return model_call
+    
+    def restore(self, seqs, align=False, **kwargs):
+        """Restore masked residues in antibody sequences."""
+        if isinstance(seqs, str):
+            seqs = [seqs]
+        
+        n_seqs = len(seqs)
+        
+        if align:
+            # For now, skip alignment as it requires ANARCI
+            print("WARNING: Alignment not implemented, skipping...")
+            pass
+        
+        # Tokenize sequences using original interface
+        tokens = self.tokenizer(seqs, pad=True, w_extra_tkns=False, device=self.used_device)
+        
+        # Get predictions for amino acids (indices 1-20)
+        predictions = self.AbLang(tokens)[:,:,1:21]
+        
+        # Find predicted tokens and replace mask tokens
+        predicted_tokens = torch.max(predictions, -1).indices + 1
+        restored_tokens = torch.where(tokens==23, predicted_tokens, tokens)
+        
+        # Decode back to sequences using original tokenizer
+        restored_seqs = self.tokenizer(restored_tokens, mode="decode")
+        
+        # Handle paired sequences format
+        if n_seqs < len(restored_seqs):
+            restored_seqs = [f"{h}|{l}".replace('-','') for h,l in zip(restored_seqs[:n_seqs], restored_seqs[n_seqs:])]
+            seqs = [f"{h}|{l}" for h,l in zip(seqs[:n_seqs], seqs[n_seqs:])]
+        
+        # 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)]])
 
 def add_angle_brackets(seq):
     # Assumes input is 'VH|VL' or 'VH|' or '|VL'