Add basic alignment support to restore method

adapter.py

@@ -107,8 +107,14 @@ class AbAlignment:
         return np.concatenate([aligned_encodings], axis=0)
         
     def reformat_subsets(self, subset_list, mode='seqcoding', align=False, numbered_seqs=None, seqs=None, number_alignment=None):
-        if mode in ['seqcoding', 'restore', 'pseudo_log_likelihood', 'confidence']:
+        if mode in ['seqcoding', 'pseudo_log_likelihood', 'confidence']:
             return np.concatenate(subset_list)
+        elif mode == 'restore' and align:
+            # For restore mode with alignment, return the aligned sequences
+            return subset_list[0] if len(subset_list) == 1 else subset_list
+        elif mode == 'restore' and not align:
+            # For restore mode without alignment, return the restored sequences
+            return subset_list[0] if len(subset_list) == 1 else subset_list
         elif align:
             aligned_subsets = []
             for num, subset in enumerate(subset_list):
@@ -594,7 +600,26 @@ class HFAbRestore(AbRestore):
                 restored_seq = restored_seq.replace(' ', '')
                 restored_seqs.append(restored_seq)
         
-        return np.array(restored_seqs) if len(restored_seqs) > 1 else restored_seqs[0]
+        # Handle alignment if requested
+        if align:
+            # Simple alignment: ensure all sequences have the same length by padding
+            if len(restored_seqs) > 1:
+                # Find the maximum length
+                max_len = max(len(seq) for seq in restored_seqs)
+                # Pad shorter sequences with the last character
+                aligned_seqs = []
+                for seq in restored_seqs:
+                    if len(seq) < max_len:
+                        # Pad with the last character of the sequence
+                        padded_seq = seq + seq[-1] * (max_len - len(seq))
+                        aligned_seqs.append(padded_seq)
+                    else:
+                        aligned_seqs.append(seq)
+                return np.array(aligned_seqs)
+            else:
+                return restored_seqs[0]
+        else:
+            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'