Implement proper alignment functionality using ANARCI

adapter.py

@@ -568,46 +568,84 @@ class HFAbRestore(AbRestore):
         n_seqs = len(seqs)
         
         if align:
-            # For now, skip alignment as it requires ANARCI
-            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
+            # Implement alignment using ANARCI to create spread sequences
+            print("DEBUG: Using alignment to create spread sequences...")
+            seqs = self._sequence_aligning(seqs)
+            nr_seqs = len(seqs)//self.spread
+            
+            tokens = self.tokenizer(seqs, pad=True, w_extra_tkns=False, device=self.used_device)          
+            predictions = self.AbLang(tokens)[:,:,1:21]
+
+            # Reshape
+            tokens = tokens.reshape(nr_seqs, self.spread, -1)
+            predictions = predictions.reshape(nr_seqs, self.spread, -1, 20)
+            seqs = seqs.reshape(nr_seqs, -1)
+
+            # Find index of best predictions
+            best_seq_idx = torch.argmax(torch.max(predictions, -1).values[:,:,1:2].mean(2), -1)
+
+            # Select best predictions           
+            tokens = tokens.gather(1, best_seq_idx.view(-1, 1).unsqueeze(1).repeat(1, 1, tokens.shape[-1])).squeeze(1)
+            predictions = predictions[range(predictions.shape[0]), best_seq_idx]
+            seqs = np.take_along_axis(seqs, best_seq_idx.view(-1, 1).cpu().numpy(), axis=1)
+        else:
+            print(f"DEBUG: Processing sequences without alignment: {seqs}")
+            tokens = self.tokenizer(seqs, pad=True, w_extra_tkns=False, device=self.used_device)
+            predictions = self.AbLang(tokens)[:,:,1:21]
+
         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):
             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
         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 _sequence_aligning(self, seqs):
+        """Create spread sequences using ANARCI alignment."""
+        tmp_seqs = [pairs.replace(">", "").replace("<", "").split("|") for pairs in seqs]
+        
+        spread_heavy = [f"<{seq}>" for seq in self._create_spread_of_sequences(tmp_seqs, chain = 'H')]
+        spread_light = [f"<{seq}>" for seq in self._create_spread_of_sequences(tmp_seqs, chain = 'L')]
+        
+        return np.concatenate([np.array(spread_heavy),np.array(spread_light)])
+    
+    def _create_spread_of_sequences(self, seqs, chain = 'H'):
+        """Create spread sequences using ANARCI."""
+        import pandas as pd
+        import anarci
+        
+        chain_idx = 0 if chain == 'H' else 1
+        numbered_seqs = anarci.run_anarci(
+            pd.DataFrame([seq[chain_idx].replace('*', 'X') for seq in seqs]).reset_index().values.tolist(), 
+            ncpu=self.ncpu, 
+            scheme='imgt',
+            allowed_species=['human', 'mouse'],
+        )
+        
+        anarci_data = pd.DataFrame(
+            [str(anarci[0][0]) if anarci else 'ANARCI_error' for anarci in numbered_seqs[1]], 
+            columns=['anarci']
+        ).astype('<U90')
+        
+        max_position = 128 if chain == 'H' else 127
+        
+        from extra_utils import get_sequences_from_anarci
+        seqs = anarci_data.apply(
+            lambda x: get_sequences_from_anarci(
+                x.anarci, 
+                max_position, 
+                self.spread
+            ), axis=1, result_type='expand'
+        ).to_numpy().reshape(-1)
+        
+        return seqs
 
 def add_angle_brackets(seq):
     # Assumes input is 'VH|VL' or 'VH|' or '|VL'