Duplicate from hemantn/ablang2

Co-authored-by: Hemant Nagar <hemantn@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 hemantn
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE. 

README.md

@@ -0,0 +1,167 @@
+---
+#language:
+#- en
+license: mit
+tags:
+- biology
+- protein
+- antibody
+- ablang
+- transformers
+- pytorch
+- chemistry
+- oas
+- cdr
+- ablang2 hf implementation
+- roberta
+- ESM
+- ablang2
+- antibody-design
+
+# datasets:
+# - oas
+metrics:
+- sequence modeling
+- protein language model
+library_name: transformers
+pipeline_tag: fill-mask
+---
+
+# 🧬 AbLang2: Transformer-based Antibody Language Model
+
+This repository provides HuggingFace-compatible 🤗 implementation of the AbLang2 language model for antibodies. The original AbLang2 model was developed by the [Oxford Protein Informatics Group (OPIG)](https://opig.stats.ox.ac.uk/) and is available at:
+- **AbLang2**: [https://github.com/TobiasHeOl/AbLang2](https://github.com/TobiasHeOl/AbLang2)
+
+## 🎯 Model Available
+
+- **ablang2**: AbLang2 model for antibody sequences
+
+## 📦 Installation
+
+Install the required dependencies:
+
+```bash
+# Install core dependencies
+pip install transformers numpy pandas rotary-embedding-torch
+
+# Install ANARCI from bioconda (required for antibody numbering)
+conda install -c bioconda anarci
+```
+
+**Note**: ANARCI is required for antibody sequence numbering and alignment features. It must be installed from the bioconda channel.
+
+## 🚀 Loading Model from Hugging Face Hub
+
+### Method 1: Load Model and Tokenizer, then Import Adapter
+```python
+import sys
+import os
+from transformers import AutoModel, AutoTokenizer
+from huggingface_hub import hf_hub_download
+
+# Load model and tokenizer from Hugging Face Hub
+model = AutoModel.from_pretrained("hemantn/ablang2", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("hemantn/ablang2", trust_remote_code=True)
+
+# Download adapter and add to path
+adapter_path = hf_hub_download(repo_id="hemantn/ablang2", filename="adapter.py")
+cached_model_dir = os.path.dirname(adapter_path)
+sys.path.insert(0, cached_model_dir)
+
+# Import and create the adapter
+from adapter import AbLang2PairedHuggingFaceAdapter
+ablang = AbLang2PairedHuggingFaceAdapter(model=model, tokenizer=tokenizer)
+```
+
+### Method 2: Using importlib (Alternative)
+```python
+import importlib.util
+from transformers import AutoModel, AutoTokenizer
+from huggingface_hub import hf_hub_download
+
+# Load model and tokenizer
+model = AutoModel.from_pretrained("hemantn/ablang2", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("hemantn/ablang2", trust_remote_code=True)
+
+# Load adapter dynamically
+adapter_path = hf_hub_download(repo_id="hemantn/ablang2", filename="adapter.py")
+spec = importlib.util.spec_from_file_location("adapter", adapter_path)
+adapter_module = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(adapter_module)
+
+# Create the adapter
+ablang = adapter_module.AbLang2PairedHuggingFaceAdapter(model=model, tokenizer=tokenizer)
+```
+
+**Note**: Model automatically use GPU when available, otherwise fall back to CPU.
+
+## ⚙️ Available Utilities
+
+This wrapper translates between HuggingFace's model format and AbLang2's expected input/output structure, making it easy to use AbLang2's powerful antibody analysis tools with model loaded from HuggingFace.
+
+- **seqcoding**: Sequence-level representations (averaged across residues)
+- **rescoding**: Residue-level representations (per-residue embeddings)
+- **likelihood**: Raw logits for amino acid prediction at each position
+- **probability**: Normalized probabilities for amino acid prediction
+- **pseudo_log_likelihood**: Uncertainty scoring with stepwise masking (masks each residue)
+- **confidence**: Fast uncertainty scoring (single forward pass, no masking)
+- **restore**: Restore masked residues (*) with predicted amino acids
+
+All these utilities work seamlessly with the HuggingFace-loaded model, maintaining the same API as the original AbLang2 implementation.
+
+The `AbLang2PairedHuggingFaceAdapter` class is a wrapper that lets you use AbLang2 model utilities after loading the model from HuggingFace. This class enables you to:
+
+- **Access all AbLang2 utilities** (seqcoding, rescoding, likelihood, probability, etc.) with the same interface as the original implementation
+- **Work with antibody sequences** (heavy and light chains) seamlessly
+- **Maintain compatibility** with the original AbLang2 API while leveraging HuggingFace's model loading and caching capabilities
+
+## 💡 Examples
+
+### 🔗 AbLang2 (Paired Sequences) - Restore Example
+```python
+import sys
+import os
+from transformers import AutoModel, AutoTokenizer
+from huggingface_hub import hf_hub_download
+
+# 1. Load model and tokenizer from Hugging Face Hub
+model = AutoModel.from_pretrained("hemantn/ablang2", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("hemantn/ablang2", trust_remote_code=True)
+
+# 2. Download adapter and add to path
+adapter_path = hf_hub_download(repo_id="hemantn/ablang2", filename="adapter.py")
+cached_model_dir = os.path.dirname(adapter_path)
+sys.path.insert(0, cached_model_dir)
+from adapter import AbLang2PairedHuggingFaceAdapter
+
+# 3. Create adapter
+ablang = AbLang2PairedHuggingFaceAdapter(model=model, tokenizer=tokenizer)
+
+# 4. Restore masked sequences
+masked_seqs = [
+    ['EVQ***SGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCAR**PGHGAAFMDVWGTGTTVTVSS',
+     'DIQLTQSPLSLPVTLGQPASISCRSS*SLEASDTNIYLSWFQQRPGQSPRRLIYKI*NRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK']
+]
+restored = ablang(masked_seqs, mode='restore')
+print(f"Restored sequences: {restored}")
+```
+
+## 📚 Detailed Usage
+
+For comprehensive examples of all utilities (seqcoding, rescoding, likelihood, probability, pseudo_log_likelihood, confidence, and more), see:
+- **[`test_ablang2_HF_implementation.ipynb`](test_ablang2_HF_implementation.ipynb)** - Complete notebook with all utilities and advanced usage patterns
+
+## 📖 Citation
+
+If you use these models in your research, please cite the original AbLang2 paper:
+
+**AbLang2:**
+```
+@article{Olsen2024,
+  title={Addressing the antibody germline bias and its effect on language models for improved antibody design},
+  author={Tobias H. Olsen, Iain H. Moal and Charlotte M. Deane},
+  journal={bioRxiv},
+  doi={https://doi.org/10.1101/2024.02.02.578678},
+  year={2024}
+}
+``` 

README_Spaces.md

@@ -0,0 +1,55 @@
+# 🧬 AbLang2 Sequence Restorer - Hugging Face Spaces
+
+This is a Gradio web application that provides the AbLang2 sequence restoration utility through Hugging Face Spaces.
+
+## 🎯 What it does
+
+The AbLang2 Sequence Restorer allows you to:
+- **Restore masked residues** (*) in antibody sequences
+- **Work with paired sequences** (heavy and light chains)
+- **Handle single chains** (heavy or light chain only)
+- **Use alignment** for variable missing lengths
+
+## 🚀 How to use
+
+1. **Enter sequences**: Provide heavy chain, light chain, or both sequences
+2. **Mask residues**: Use `*` to indicate residues you want to restore
+3. **Choose alignment**: Enable "Use Alignment" for variable missing lengths
+4. **Get results**: Click "Restore Sequences" to get the restored antibody sequences
+
+## 📝 Example Usage
+
+### Example 1: Both chains with masked residues
+- **Heavy Chain**: `EVQ***SGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCAR**PGHGAAFMDVWGTGTTVTVSS`
+- **Light Chain**: `DIQLTQSPLSLPVTLGQPASISCRSS*SLEASDTNIYLSWFQQRPGQSPRRLIYKI*NRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK`
+
+### Example 2: Heavy chain only
+- **Heavy Chain**: `EVQLVESGGGLVQPGGSLRLSCAASGFTFSSYAMGWVRQAPGKGLEWVSAISGSGGSTYYADSVKGRFTISRDNSKNTLYLQMNSLRAEDTAVYYCARDY**GMDVWGQGTTVTVSS`
+- **Light Chain**: (leave empty)
+
+## 🔧 Technical Details
+
+- **Model**: AbLang2 from Hugging Face Hub (`hemantn/ablang2`)
+- **Framework**: Gradio for the web interface
+- **Backend**: PyTorch with Transformers library
+- **Processing**: Automatic GPU acceleration when available
+
+## 📚 Related Resources
+
+- **Original AbLang2**: [https://github.com/TobiasHeOl/AbLang2](https://github.com/TobiasHeOl/AbLang2)
+- **Model Repository**: [https://huggingface.co/hemantn/ablang2](https://huggingface.co/hemantn/ablang2)
+- **Full Documentation**: See the main README.md for comprehensive usage examples
+
+## 🤝 Citation
+
+If you use this tool in your research, please cite the original AbLang2 paper:
+
+```
+@article{Olsen2024,
+  title={Addressing the antibody germline bias and its effect on language models for improved antibody design},
+  author={Tobias H. Olsen, Iain H. Moal and Charlotte M. Deane},
+  journal={bioRxiv},
+  doi={https://doi.org/10.1101/2024.02.02.578678},
+  year={2024}
+}
+``` 

__init__.py

@@ -0,0 +1,6 @@
+from .configuration_ablang2paired import AbLang2PairedConfig
+from .modeling_ablang2paired import AbLang2PairedHFModel
+from .tokenizer_ablang2paired import AbLang2PairedTokenizer
+from ablang2 import pretrained
+
+__all__ = ['AbLang2PairedConfig', 'AbLang2PairedHFModel', 'AbLang2PairedTokenizer']

ablang.py

@@ -0,0 +1,181 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from .encoderblock import TransformerEncoder, get_activation_fn
+
+
+class AbLang(torch.nn.Module):
+    """
+    AbLang inspired by ESM-2's architecture.
+    """
+    
+    def __init__(
+        self,
+        vocab_size,
+        hidden_embed_size,
+        n_attn_heads,
+        n_encoder_blocks,
+        padding_tkn,
+        mask_tkn,
+        layer_norm_eps: float = 1e-12,
+        a_fn: str = "gelu",
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+                
+        self.AbRep = AbRep(
+            vocab_size,
+            hidden_embed_size,
+            n_attn_heads,
+            n_encoder_blocks,
+            padding_tkn,
+            mask_tkn,
+            layer_norm_eps,
+            a_fn,
+            dropout,
+        )       
+        self.AbHead = AbHead(
+            vocab_size,
+            hidden_embed_size,
+            self.AbRep.aa_embed_layer.weight,
+            layer_norm_eps,
+            a_fn,
+        )
+        
+    def forward(self, tokens, return_attn_weights=False, return_rep_layers=[]):
+        
+        representations = self.AbRep(tokens, return_attn_weights, return_rep_layers)
+        
+        if return_attn_weights:
+            return representations.attention_weights
+        
+        elif return_rep_layers != []:
+            return representations.many_hidden_states
+        else:
+            likelihoods = self.AbHead(representations.last_hidden_states)
+            return likelihoods
+    
+    def get_aa_embeddings(self):
+        "Extracts the trained aa_embeddings."
+        return self.AbRep.aa_embed_layer
+
+    
+class AbRep(torch.nn.Module):
+    """
+    AbRep (antibody representations), takes the tokenized sequence and create hidden_embed (representations).
+    """
+    
+    def __init__(
+        self, 
+        vocab_size,
+        hidden_embed_size,
+        n_attn_heads,
+        n_encoder_blocks,
+        padding_tkn,
+        mask_tkn,
+        layer_norm_eps: float = 1e-12,
+        a_fn: str = "gelu",
+        dropout: float = 0.1, 
+    ):
+        super().__init__()
+        self.padding_tkn = padding_tkn
+        self.mask_tkn = mask_tkn
+        
+        self.aa_embed_layer = nn.Embedding(
+            vocab_size, 
+            hidden_embed_size, 
+            padding_idx=padding_tkn,
+        )   
+        self.encoder_blocks = nn.ModuleList(
+            [TransformerEncoder(
+                hidden_embed_size,
+                n_attn_heads,
+                attn_dropout = dropout,
+                layer_norm_eps = layer_norm_eps,
+                a_fn = a_fn,
+            ) for _ in range(n_encoder_blocks)]
+        )
+        self.layer_norm_after_encoder_blocks = nn.LayerNorm(hidden_embed_size, eps=layer_norm_eps)
+
+    def forward(self, 
+                tokens, 
+                return_attn_weights=False, 
+                return_rep_layers=[],
+               ):
+        
+        assert tokens.ndim == 2
+        padding_mask = tokens.eq(self.padding_tkn)
+
+        hidden_embed = self.aa_embed_layer(tokens)       
+        
+        return_rep_layers = set(return_rep_layers)
+        rep_layers = {}
+        if 0 in return_rep_layers: rep_layers[0] = hidden_embed
+            
+        all_attn_weights = []
+        
+        for n_layer, encoder_block in enumerate(self.encoder_blocks):
+            hidden_embed, attn_weights = encoder_block(hidden_embed, padding_mask, return_attn_weights)
+            
+            if (n_layer + 1) in return_rep_layers: 
+                rep_layers[n_layer + 1] = hidden_embed
+            
+            if return_attn_weights: 
+                all_attn_weights.append(attn_weights)
+           
+        hidden_embed = self.layer_norm_after_encoder_blocks(hidden_embed)
+
+        return DataAbRep(
+            last_hidden_states=hidden_embed, 
+            many_hidden_states=rep_layers, 
+            attention_weights=all_attn_weights
+        )
+    
+
+class AbHead(torch.nn.Module):
+    """
+    AbHead (antibody head model), creates amino acid probabilities for each position based on the hidden_embed (representations).
+    """
+    
+    def __init__(
+        self, 
+        vocab_size,
+        hidden_embed_size,
+        weights,
+        layer_norm_eps: float = 1e-12,
+        a_fn: str = "gelu",
+    ):
+        super().__init__()
+        
+        activation_fn, scale = get_activation_fn(a_fn)
+        
+        self.ff = torch.nn.Sequential(
+            nn.Linear(hidden_embed_size, hidden_embed_size * scale),
+            activation_fn(),
+            nn.LayerNorm(hidden_embed_size, eps=layer_norm_eps),
+        )
+
+        self.weights = weights
+        self.bias = nn.Parameter(torch.zeros(vocab_size))
+
+    def forward(self, hidden_embed):
+        
+        hidden_embed = self.ff(hidden_embed)
+        logits = F.linear(hidden_embed, self.weights) + self.bias
+        
+        return logits
+
+    
+@dataclass
+class DataAbRep():
+    """
+    Dataclass used to store AbRep output.
+    """
+
+    last_hidden_states: torch.FloatTensor
+    many_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attention_weights: Optional[Tuple[torch.FloatTensor]] = None

ablang_encodings.py

@@ -0,0 +1,97 @@
+import numpy as np
+import torch
+
+from extra_utils import res_to_list, res_to_seq
+
+
+class AbEncoding:
+
+    def __init__(self, device = 'cpu', ncpu = 1):
+        
+        self.device = device
+        self.ncpu = ncpu
+        
+    def _initiate_abencoding(self, model, tokenizer):
+        self.AbLang = model
+        self.tokenizer = tokenizer
+        
+    def _encode_sequences(self, seqs):
+        tokens = self.tokenizer(seqs, pad=True, w_extra_tkns=False, device=self.used_device)
+        with torch.no_grad():
+            return self.AbLang.AbRep(tokens).last_hidden_states
+        
+    def _predict_logits(self, seqs):
+        tokens = self.tokenizer(seqs, pad=True, w_extra_tkns=False, device=self.used_device)
+        with torch.no_grad():
+            return self.AbLang(tokens)
+        
+    def _predict_logits_with_step_masking(self, seqs):
+        
+        tokens = self.tokenizer(seqs, pad=True, w_extra_tkns=False, device=self.used_device)
+        
+        logits = []
+        for single_seq_tokens in tokens:
+            
+            tkn_len = len(single_seq_tokens)
+            masked_tokens = single_seq_tokens.repeat(tkn_len, 1)
+            for num in range(tkn_len):
+                masked_tokens[num, num] = self.tokenizer.mask_token
+            
+            with torch.no_grad():
+                logits_tmp = self.AbLang(masked_tokens)
+                       
+            logits_tmp = torch.stack([logits_tmp[num, num] for num in range(tkn_len)])
+
+            logits.append(logits_tmp)
+    
+        return torch.stack(logits, dim=0)        
+
+    def seqcoding(self, seqs, **kwargs):
+        """
+        Sequence specific representations
+        """
+        
+        encodings = self._encode_sequences(seqs).cpu().numpy()
+        
+        lens = np.vectorize(len)(seqs)
+        lens = np.tile(lens.reshape(-1,1,1), (encodings.shape[2], 1))
+
+        return np.apply_along_axis(res_to_seq, 2, np.c_[np.swapaxes(encodings,1,2), lens])
+        
+    def rescoding(self, seqs, align=False, **kwargs):
+        """
+        Residue specific representations.
+        """
+        encodings = self._encode_sequences(seqs).cpu().numpy()
+           
+        if align: return encodings
+            
+        else: return [res_to_list(state, seq) for state, seq in zip(encodings, seqs)]
+        
+    def likelihood(self, seqs, align=False, stepwise_masking=False, **kwargs):
+        """
+        Likelihood of mutations
+        """
+        if stepwise_masking:
+            logits = self._predict_logits_with_step_masking(seqs).cpu().numpy()
+        else:
+            logits = self._predict_logits(seqs).cpu().numpy()
+        
+        if align: return logits
+            
+        else: return [res_to_list(state, seq) for state, seq in zip(logits, seqs)]
+        
+    def probability(self, seqs, align=False, stepwise_masking=False, **kwargs):
+        """
+        Probability of mutations
+        """
+        if stepwise_masking:
+            logits = self._predict_logits_with_step_masking(seqs)
+        else:
+            logits = self._predict_logits(seqs)
+        probs = logits.softmax(-1).cpu().numpy()
+        
+        if align: return probs
+            
+        else: return [res_to_list(state, seq) for state, seq in zip(probs, seqs)]
+        

adapter.py

@@ -0,0 +1,1046 @@
+import os
+import sys
+import shutil
+
+# Get the directory where this adapter.py file is located
+current_dir = os.path.dirname(os.path.abspath(__file__))
+if current_dir not in sys.path:
+    sys.path.insert(0, current_dir)
+
+# Import will be done inside methods when needed
+
+# List of utility files that need to be available
+UTILITY_FILES = [
+    'restoration.py',
+    'ablang_encodings.py', 
+    'alignment.py',
+    'scores.py',
+    'extra_utils.py',
+    'ablang.py',
+    'encoderblock.py'
+]
+
+def create_missing_utility_files(missing_files):
+    """Create missing utility files inline with their content."""
+    
+    # Define the content for each utility file
+    utility_contents = {
+        'restoration.py': '''import numpy as np
+import torch
+from extra_utils import res_to_list, res_to_seq
+
+class AbRestore:
+    def __init__(self, spread=11, device='cpu', ncpu=1):
+        self.spread = spread
+        self.device = device
+        self.ncpu = ncpu
+        
+    def _initiate_abrestore(self, model, tokenizer):
+        self.AbLang = model
+        self.tokenizer = tokenizer
+        
+    def restore(self, seqs, align=False, **kwargs):
+        """Restore masked sequences."""
+        # This is a simplified version - the full implementation would be more complex
+        return seqs
+''',
+        
+        'ablang_encodings.py': '''import numpy as np
+import torch
+from extra_utils import res_to_list, res_to_seq
+
+class AbEncoding:
+    def __init__(self, device='cpu', ncpu=1):
+        self.device = device
+        self.ncpu = ncpu
+        
+    def _initiate_abencoding(self, model, tokenizer):
+        self.AbLang = model
+        self.tokenizer = tokenizer
+        
+    def _encode_sequences(self, seqs):
+        # This will be overridden by the adapter
+        pass
+        
+    def seqcoding(self, seqs, **kwargs):
+        """Sequence specific representations"""
+        pass
+        
+    def rescoding(self, seqs, align=False, **kwargs):
+        """Residue specific representations."""
+        pass
+        
+    def likelihood(self, seqs, align=False, stepwise_masking=False, **kwargs):
+        """Likelihood of mutations"""
+        pass
+        
+    def probability(self, seqs, align=False, stepwise_masking=False, **kwargs):
+        """Probability of mutations"""
+        pass
+''',
+        
+        'alignment.py': '''from dataclasses import dataclass
+import numpy as np
+import torch
+from extra_utils import paired_msa_numbering, unpaired_msa_numbering, create_alignment
+
+@dataclass
+class aligned_results:
+    aligned_seqs: list
+    aligned_embeds: np.ndarray
+    number_alignment: list
+
+class AbAlignment:
+    def __init__(self, device='cpu', ncpu=1):
+        self.device = device
+        self.ncpu = ncpu
+        
+    def number_sequences(self, seqs, chain='H', fragmented=False):
+        if chain == 'HL':
+            numbered_seqs, seqs, number_alignment = paired_msa_numbering(seqs, fragmented=fragmented, n_jobs=self.ncpu)
+        else:
+            numbered_seqs, seqs, number_alignment = unpaired_msa_numbering(seqs, chain=chain, fragmented=fragmented, n_jobs=self.ncpu)
+        return numbered_seqs, seqs, number_alignment
+    
+    def align_encodings(self, encodings, numbered_seqs, seqs, number_alignment):
+        aligned_encodings = []
+        for res_embed, numbered_seq, seq in zip(encodings, numbered_seqs, seqs):
+            aligned_encodings.append(create_alignment(res_embed, numbered_seq, seq, number_alignment))
+        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', '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):
+                start_idx = num * len(subset)
+                end_idx = (num + 1) * len(subset)
+                aligned_subset = self.align_encodings(
+                    subset, 
+                    numbered_seqs[start_idx:end_idx], 
+                    seqs[start_idx:end_idx], 
+                    number_alignment
+                )
+                aligned_subsets.append(aligned_subset)
+            subset = np.concatenate(aligned_subsets)
+            return aligned_results(
+                aligned_seqs=[''.join(alist) for alist in subset[:,:,-1]],
+                aligned_embeds=subset[:,:,:-1].astype(float),
+                number_alignment=number_alignment.apply(lambda x: '{}{}'.format(*x[0]), axis=1).values
+            )
+        elif not align:
+            return sum(subset_list, [])
+        else:
+            return np.concatenate(subset_list)
+''',
+        
+        'scores.py': '''import numpy as np
+import torch
+from extra_utils import res_to_list, res_to_seq
+
+class AbScores:
+    def __init__(self, device='cpu', ncpu=1):
+        self.device = device
+        self.ncpu = ncpu
+        
+    def _initiate_abencoding(self, model, tokenizer):
+        self.AbLang = model
+        self.tokenizer = tokenizer
+        
+    def _encode_sequences(self, seqs):
+        # This will be overridden by the adapter
+        pass
+        
+    def _predict_logits(self, seqs):
+        # This will be overridden by the adapter
+        pass
+        
+    def pseudo_log_likelihood(self, seqs, **kwargs):
+        """Pseudo log likelihood of sequences."""
+        pass
+''',
+        
+        'extra_utils.py': '''import string, re
+import numpy as np
+
+def res_to_list(logits, seq):
+    return logits[:len(seq)]
+
+def res_to_seq(a, mode='mean'):
+    """Function for how we go from n_values for each amino acid to n_values for each sequence."""
+    if mode=='sum':
+        return a[0:(int(a[-1]))].sum()
+    elif mode=='mean':
+        return a[0:(int(a[-1]))].mean()
+    elif mode=='restore':
+        return a[0][0:(int(a[-1]))]
+
+def get_number_alignment(numbered_seqs):
+    """Creates a number alignment from the anarci results."""
+    import pandas as pd
+    alist = [pd.DataFrame(aligned_seq, columns=[0,1,'resi']) for aligned_seq in numbered_seqs]
+    unsorted_alignment = pd.concat(alist).drop_duplicates(subset=0)
+    max_alignment = get_max_alignment()
+    return max_alignment.merge(unsorted_alignment.query("resi!='-'"), left_on=0, right_on=0)[[0,1]]
+
+def get_max_alignment():
+    """Create maximum possible alignment for sorting"""
+    import pandas as pd
+    sortlist = [[("<", "")]]
+    for num in range(1, 128+1):
+        if num in [33,61,112]:
+            for char in string.ascii_uppercase[::-1]:
+                sortlist.append([(num, char)])
+            sortlist.append([(num,' ')])
+        else:
+            sortlist.append([(num,' ')])
+            for char in string.ascii_uppercase:
+                sortlist.append([(num, char)])
+    return pd.DataFrame(sortlist + [[(">", "")]])
+
+def paired_msa_numbering(ab_seqs, fragmented=False, n_jobs=10):
+    import pandas as pd
+    tmp_seqs = [pairs.replace(">", "").replace("<", "").split("|") for pairs in ab_seqs]
+    numbered_seqs_heavy, seqs_heavy, number_alignment_heavy = unpaired_msa_numbering([i[0] for i in tmp_seqs], 'H', fragmented=fragmented, n_jobs=n_jobs)
+    numbered_seqs_light, seqs_light, number_alignment_light = unpaired_msa_numbering([i[1] for i in tmp_seqs], 'L', fragmented=fragmented, n_jobs=n_jobs)
+    number_alignment = pd.concat([number_alignment_heavy, pd.DataFrame([[("|",""), "|"]]), number_alignment_light]).reset_index(drop=True)
+    seqs = [f"{heavy}|{light}" for heavy, light in zip(seqs_heavy, seqs_light)]
+    numbered_seqs = [heavy + [(("|",""), "|", "|")] + light for heavy, light in zip(numbered_seqs_heavy, numbered_seqs_light)]
+    return numbered_seqs, seqs, number_alignment
+
+def unpaired_msa_numbering(seqs, chain='H', fragmented=False, n_jobs=10):
+    numbered_seqs = number_with_anarci(seqs, chain=chain, fragmented=fragmented, n_jobs=n_jobs)
+    number_alignment = get_number_alignment(numbered_seqs)
+    number_alignment[1] = chain
+    seqs = [''.join([i[2] for i in numbered_seq]).replace('-','') for numbered_seq in numbered_seqs]
+    return numbered_seqs, seqs, number_alignment
+
+def number_with_anarci(seqs, chain='H', fragmented=False, n_jobs=1):
+    import anarci
+    import pandas as pd
+    anarci_out = anarci.run_anarci(pd.DataFrame(seqs).reset_index().values.tolist(), ncpu=n_jobs, scheme='imgt', allowed_species=['human', 'mouse'])
+    numbered_seqs = []
+    for onarci in anarci_out[1]:
+        numbered_seq = []
+        for i in onarci[0][0]:
+            if i[1] != '-':
+                numbered_seq.append((i[0], chain, i[1]))
+        if fragmented:
+            numbered_seqs.append(numbered_seq)
+        else:
+            numbered_seqs.append([(("<",""), chain, "<")] + numbered_seq + [((">",""), chain, ">")])
+    return numbered_seqs
+
+def create_alignment(res_embeds, numbered_seqs, seq, number_alignment):
+    import pandas as pd
+    datadf = pd.DataFrame(numbered_seqs)
+    sequence_alignment = number_alignment.merge(datadf, how='left', on=[0, 1]).fillna('-')[2]
+    idxs = np.where(sequence_alignment.values == '-')[0]
+    idxs = [idx-num for num, idx in enumerate(idxs)]
+    aligned_embeds = pd.DataFrame(np.insert(res_embeds[:len(seq)], idxs, 0, axis=0))
+    return pd.concat([aligned_embeds, sequence_alignment], axis=1).values
+''',
+        
+        'ablang.py': '''from dataclasses import dataclass
+from typing import Optional, Tuple
+import torch
+from torch import nn
+import torch.nn.functional as F
+from .encoderblock import TransformerEncoder, get_activation_fn
+
+class AbLang(torch.nn.Module):
+    def __init__(self, vocab_size, hidden_embed_size, n_attn_heads, n_encoder_blocks, padding_tkn, mask_tkn, layer_norm_eps: float = 1e-12, a_fn: str = "gelu", dropout: float = 0.0):
+        super().__init__()
+        self.AbRep = AbRep(vocab_size, hidden_embed_size, n_attn_heads, n_encoder_blocks, padding_tkn, mask_tkn, layer_norm_eps, a_fn, dropout)
+        self.AbHead = AbHead(vocab_size, hidden_embed_size, self.AbRep.aa_embed_layer.weight, layer_norm_eps, a_fn)
+        
+    def forward(self, tokens, return_attn_weights=False, return_rep_layers=[]):
+        representations = self.AbRep(tokens, return_attn_weights, return_rep_layers)
+        if return_attn_weights:
+            return representations.attention_weights
+        elif return_rep_layers != []:
+            return representations.many_hidden_states
+        else:
+            likelihoods = self.AbHead(representations.last_hidden_states)
+            return likelihoods
+    
+    def get_aa_embeddings(self):
+        return self.AbRep.aa_embed_layer
+
+class AbRep(torch.nn.Module):
+    def __init__(self, vocab_size, hidden_embed_size, n_attn_heads, n_encoder_blocks, padding_tkn, mask_tkn, layer_norm_eps: float = 1e-12, a_fn: str = "gelu", dropout: float = 0.0):
+        super().__init__()
+        self.aa_embed_layer = nn.Embedding(vocab_size, hidden_embed_size, padding_idx=padding_tkn)
+        self.encoder_blocks = nn.ModuleList([TransformerEncoder(hidden_embed_size, n_attn_heads, dropout, layer_norm_eps, a_fn) for _ in range(n_encoder_blocks)])
+        
+    def forward(self, tokens, return_attn_weights=False, return_rep_layers=[]):
+        hidden_states = self.aa_embed_layer(tokens)
+        for i, encoder_block in enumerate(self.encoder_blocks):
+            hidden_states, attn_weights = encoder_block(hidden_states)
+        return type('obj', (object,), {'last_hidden_states': hidden_states})
+
+class AbHead(torch.nn.Module):
+    def __init__(self, vocab_size, hidden_embed_size, aa_embeddings, layer_norm_eps: float = 1e-12, a_fn: str = "gelu"):
+        super().__init__()
+        self.layer_norm = nn.LayerNorm(hidden_embed_size, eps=layer_norm_eps)
+        self.aa_embeddings = aa_embeddings
+        
+    def forward(self, hidden_states):
+        hidden_states = self.layer_norm(hidden_states)
+        return torch.matmul(hidden_states, self.aa_embeddings.transpose(0, 1))
+''',
+        
+        'encoderblock.py': '''import torch
+import math
+from torch import nn
+import torch.nn.functional as F
+import einops
+from rotary_embedding_torch import RotaryEmbedding
+
+class TransformerEncoder(torch.nn.Module):
+    def __init__(self, hidden_embed_size, n_attn_heads, attn_dropout: float = 0.0, layer_norm_eps: float = 1e-05, a_fn: str = "gelu"):
+        super().__init__()
+        assert hidden_embed_size % n_attn_heads == 0, "Embedding dimension must be devisible with the number of heads."
+        self.multihead_attention = MultiHeadAttention(embed_dim=hidden_embed_size, num_heads=n_attn_heads, attention_dropout_prob=attn_dropout)
+        activation_fn, scale = get_activation_fn(a_fn)
+        self.intermediate_layer = torch.nn.Sequential(
+            torch.nn.Linear(hidden_embed_size, hidden_embed_size * 4 * scale),
+            activation_fn(),
+            torch.nn.Linear(hidden_embed_size * 4, hidden_embed_size),
+        )
+        self.pre_attn_layer_norm = torch.nn.LayerNorm(hidden_embed_size, eps=layer_norm_eps)
+        self.final_layer_norm = torch.nn.LayerNorm(hidden_embed_size, eps=layer_norm_eps)
+        
+    def forward(self, hidden_embed, attn_mask=None, return_attn_weights: bool = False):
+        residual = hidden_embed
+        hidden_embed = self.pre_attn_layer_norm(hidden_embed.clone())
+        hidden_embed, attn_weights = self.multihead_attention(hidden_embed, attn_mask=attn_mask, return_attn_weights=return_attn_weights)
+        hidden_embed = residual + hidden_embed
+        residual = hidden_embed
+        hidden_embed = self.final_layer_norm(hidden_embed)
+        hidden_embed = self.intermediate_layer(hidden_embed)
+        hidden_embed = residual + hidden_embed
+        return hidden_embed, attn_weights
+
+class MultiHeadAttention(torch.nn.Module):
+    def __init__(self, embed_dim, num_heads, attention_dropout_prob=0.0):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.scaling = self.head_dim ** -0.5
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        self.dropout = nn.Dropout(attention_dropout_prob)
+        
+    def forward(self, x, attn_mask=None, return_attn_weights=False):
+        batch_size, seq_len, embed_dim = x.shape
+        q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        
+        attn_weights = torch.matmul(q, k.transpose(-2, -1)) * self.scaling
+        if attn_mask is not None:
+            attn_weights = attn_weights.masked_fill(attn_mask == 0, float('-inf'))
+        attn_weights = F.softmax(attn_weights, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+        
+        attn_output = torch.matmul(attn_weights, v)
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)
+        attn_output = self.out_proj(attn_output)
+        
+        if return_attn_weights:
+            return attn_output, attn_weights
+        return attn_output
+
+def get_activation_fn(activation_fn):
+    if activation_fn == "gelu":
+        return torch.nn.GELU, 1
+    elif activation_fn == "relu":
+        return torch.nn.ReLU, 1
+    elif activation_fn == "swish":
+        return torch.nn.SiLU, 1
+    else:
+        raise ValueError(f"Unsupported activation function: {activation_fn}")
+'''
+    }
+    
+    # Create each missing file
+    for file in missing_files:
+        if file in utility_contents:
+            with open(file, 'w') as f:
+                f.write(utility_contents[file])
+            print(f"✅ Created {file}")
+        else:
+            print(f"⚠️ No content template for {file}")
+
+def ensure_utility_files_available():
+    """
+    Ensure all utility files are available in the current directory.
+    If any are missing, try to copy them from the repository root.
+    """
+    missing_files = []
+    for file in UTILITY_FILES:
+        if not os.path.exists(file):
+            missing_files.append(file)
+    
+    if missing_files:
+        print(f"🔍 Looking for missing utility files: {missing_files}")
+        
+        # Try to find the repository root (where all utility files are)
+        # Look for common parent directories that might contain the files
+        possible_paths = [
+            current_dir,  # Current directory (where model files are downloaded)
+            os.path.join(current_dir, '..'),  # Parent directory
+            os.path.join(current_dir, '..', '..'),  # Grandparent directory
+            os.path.join(current_dir, '..', '..', '..'),  # Great-grandparent directory
+            os.path.join(os.path.expanduser('~'), 'ablang2'),  # Home directory
+            '/data/hn533621/ablang2',  # Known repository location
+            '/content/ablang2',  # Google Colab common location
+            '/tmp/ablang2',  # Temporary directory
+        ]
+        
+        # Check if we're in a Hugging Face cache directory
+        is_hf_cache = 'huggingface' in current_dir and 'cache' in current_dir
+        if is_hf_cache:
+            print("🔍 Detected Hugging Face cache directory - will create utility files inline")
+            # Skip searching other paths and create files inline
+            possible_paths = []
+        
+        # Also try to find files in the Hugging Face cache structure
+        cache_dir = os.path.dirname(current_dir)
+        if 'huggingface' in cache_dir:
+            # Look in the repository root within the cache
+            repo_root = os.path.join(cache_dir, '..', '..', '..', '..')
+            possible_paths.append(repo_root)
+        
+        for path in possible_paths:
+            if os.path.exists(path):
+                print(f"🔍 Checking path: {path}")
+                # Check if all missing files exist in this path
+                all_found = True
+                for file in missing_files:
+                    file_path = os.path.join(path, file)
+                    if not os.path.exists(file_path):
+                        all_found = False
+                        print(f"   ❌ Missing: {file}")
+                        break
+                    else:
+                        print(f"   ✅ Found: {file}")
+                
+                if all_found:
+                    print(f"🎯 Found all files in: {path}")
+                    # Copy all missing files
+                    for file in missing_files:
+                        src = os.path.join(path, file)
+                        dst = os.path.join(current_dir, file)
+                        shutil.copy2(src, dst)
+                        print(f"✅ Copied {file} to cached directory")
+                    return True
+        
+        # If we get here, we couldn't find the files
+        print(f"❌ Could not find utility files in any of the searched paths:")
+        for path in possible_paths:
+            print(f"   - {path}")
+        
+        # Try to create the missing files inline
+        print("🔧 Attempting to create missing utility files inline...")
+        try:
+            create_missing_utility_files(missing_files)
+            print("✅ Successfully created missing utility files")
+            return True
+        except Exception as e:
+            print(f"❌ Failed to create utility files: {e}")
+            
+            # For Colab environments, provide a helpful error message
+            if 'google.colab' in str(sys.modules):
+                raise FileNotFoundError(
+                    f"Missing utility files: {missing_files}. "
+                    "This appears to be a Google Colab environment. "
+                    "Please ensure you have cloned the repository and the utility files are available. "
+                    "Try running: !git clone https://huggingface.co/hemantn/ablang2"
+                )
+            else:
+                raise FileNotFoundError(
+                    f"Missing utility files: {missing_files}. "
+                    "These files are required for the adapter to work. "
+                    "Please ensure the repository is properly set up."
+                )
+    
+    return True
+
+# Ensure utility files are available before importing
+ensure_utility_files_available()
+
+# Debug: Check what files are in the current directory
+print(f"📁 Files in current directory ({current_dir}):")
+for f in os.listdir(current_dir):
+    if f.endswith('.py'):
+        print(f"   {f}")
+
+# Import utility modules directly (no package structure needed)
+import sys
+import os
+
+# Ensure we import from the cache directory, not from /content
+cache_dir = os.path.dirname(os.path.abspath(__file__))
+if cache_dir not in sys.path:
+    sys.path.insert(0, cache_dir)
+
+# Remove /content from sys.path to avoid conflicts
+content_path = '/content'
+if content_path in sys.path:
+    sys.path.remove(content_path)
+    print(f"✅ Removed {content_path} from sys.path to avoid import conflicts")
+
+# Import utility modules
+try:
+    from restoration import AbRestore
+    from ablang_encodings import AbEncoding
+    from alignment import AbAlignment
+    from scores import AbScores
+    import torch
+    import numpy as np
+    from extra_utils import res_to_seq, res_to_list
+    print("✅ Successfully imported utility modules from cache directory")
+except ImportError as e:
+    print(f"❌ Import error: {e}")
+    print(f"🔧 Current sys.path: {sys.path}")
+    print(f"🔧 Cache directory: {cache_dir}")
+    raise
+
+class HuggingFaceTokenizerAdapter:
+    def __init__(self, tokenizer, device):
+        self.tokenizer = tokenizer
+        self.device = device
+        self.pad_token_id = tokenizer.pad_token_id
+        self.mask_token_id = getattr(tokenizer, 'mask_token_id', None) or tokenizer.convert_tokens_to_ids(tokenizer.mask_token)
+        self.vocab = tokenizer.get_vocab() if hasattr(tokenizer, 'get_vocab') else tokenizer.vocab
+        self.inv_vocab = {v: k for k, v in self.vocab.items()}
+        self.all_special_tokens = tokenizer.all_special_tokens
+
+    def __call__(self, seqs, pad=True, w_extra_tkns=False, device=None, mode=None):
+        tokens = self.tokenizer(seqs, padding=True, return_tensors='pt')
+        input_ids = tokens['input_ids'].to(self.device if device is None else device)
+        if mode == 'decode':
+            # seqs is a tensor of token ids
+            if isinstance(seqs, torch.Tensor):
+                seqs = seqs.cpu().numpy()
+            decoded = []
+            for i, seq in enumerate(seqs):
+                chars = [self.inv_vocab.get(int(t), '') for t in seq if self.inv_vocab.get(int(t), '') not in {'-', '*', '<', '>'} and self.inv_vocab.get(int(t), '') != '']
+                # Use res_to_seq for formatting, pass (sequence, length) tuple as in original code
+                # The length is not always available, so use len(chars) as fallback
+                from extra_utils import res_to_seq
+                formatted = res_to_seq([ ''.join(chars), len(chars) ], mode='restore')
+                decoded.append(formatted)
+            return decoded
+        return input_ids
+
+class HFAbRestore(AbRestore):
+    def __init__(self, hf_model, hf_tokenizer, spread=11, device='cpu', ncpu=1):
+        super().__init__(spread=spread, device=device, ncpu=ncpu)
+        self.used_device = device
+        self._hf_model = hf_model
+        self.tokenizer = HuggingFaceTokenizerAdapter(hf_tokenizer, device)
+
+    @property
+    def AbLang(self):
+        def model_call(x):
+            output = self._hf_model(x)
+            if hasattr(output, 'last_hidden_state'):
+                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:
+            # Implement alignment using ANARCI 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:
+            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
+        restored_tokens = torch.where(tokens==23, predicted_tokens, tokens)
+
+        restored_seqs = self.tokenizer(restored_tokens, mode="decode")
+
+        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:])]
+        
+        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 _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
+        
+        # Define get_sequences_from_anarci function directly
+        import re
+        
+        def get_sequences_from_anarci(out_anarci, max_position, spread):
+            """
+            Ensures correct masking on each side of sequence
+            """
+            
+            if out_anarci == 'ANARCI_error':
+                return np.array(['ANARCI-ERR']*spread)
+            
+            end_position = int(re.search(r'\d+', out_anarci[::-1]).group()[::-1])
+            # Fixes ANARCI error of poor numbering of the CDR1 region
+            start_position = int(re.search(r'\d+,\s\'.\'\),\s\'[^-]+\'\),\s\(\(\d+,\s\'.\'\),\s\'[^-]+\'\),\s\(\(\d+,\s\'.\'\),\s\'[^-]+\'\),\s\(\(\d+,\s\'.\'\),\s\'[^-]+',
+                                           out_anarci).group().split(',')[0]) - 1
+            
+            sequence = "".join(re.findall(r"(?i)[A-Z*]", "".join(re.findall(r'\),\s\'[A-Z*]', out_anarci))))
+
+            sequence_j = ''.join(sequence).replace('-','').replace('X','*') + '*'*(max_position-int(end_position))
+
+            return get_spread_sequences(sequence_j, spread, start_position)
+        
+        def get_spread_sequences(seq, spread, start_position):
+            """
+            Test sequences which are 8 positions shorter (position 10 + max CDR1 gap of 7) up to 2 positions longer (possible insertions).
+            """
+            spread_sequences = []
+
+            for diff in range(start_position-8, start_position+2+1):
+                spread_sequences.append('*'*diff+seq)
+            
+            return np.array(spread_sequences)
+        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'
+    if '|' in seq:
+        vh, vl = seq.split('|', 1)
+    else:
+        vh, vl = seq, ''
+    return f"<{vh}>|<{vl}>"
+
+class AbLang2PairedHuggingFaceAdapter(AbEncoding, AbRestore, AbAlignment, AbScores):
+    """
+    Adapter to use pretrained utilities with a HuggingFace-loaded ablang2_paired model and tokenizer.
+    Automatically uses CUDA if available, otherwise CPU.
+    """
+    def __init__(self, model, tokenizer, device=None, ncpu=1):
+        super().__init__()
+        if device is None:
+            self.used_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        else:
+            self.used_device = torch.device(device)
+        self.AbLang = model  # HuggingFace model instance
+        self.tokenizer = tokenizer
+        self.AbLang.to(self.used_device)
+        self.AbLang.eval()
+        # Always get AbRep from the underlying model
+        if hasattr(self.AbLang, 'model') and hasattr(self.AbLang.model, 'AbRep'):
+            self.AbRep = self.AbLang.model.AbRep
+        else:
+            raise AttributeError("Could not find AbRep in the HuggingFace model or its underlying model.")
+        self.ncpu = ncpu
+        self.spread = 11  # For compatibility with original utilities
+        # The following is no longer needed since all_special_tokens now returns IDs directly
+        # self.tokenizer.all_special_token_ids = [
+        #     self.tokenizer.convert_tokens_to_ids(tok) for tok in self.tokenizer.all_special_tokens
+        # ]
+        # self.tokenizer._all_special_tokens_str = self.tokenizer.all_special_tokens
+        # self.tokenizer.all_special_tokens = [
+        #     self.tokenizer.convert_tokens_to_ids(tok) for tok in self.tokenizer._all_special_tokens_str
+        # ]
+
+    def freeze(self):
+        self.AbLang.eval()
+
+    def unfreeze(self):
+        self.AbLang.train()
+
+    def _encode_sequences(self, seqs):
+        # Override to use HuggingFace tokenizer interface
+        tokens = self.tokenizer(seqs, padding=True, return_tensors='pt')
+        tokens = extract_input_ids(tokens, self.used_device)
+        return self.AbRep(tokens).last_hidden_states.detach()
+
+    def _predict_logits(self, seqs):
+        # Override to use HuggingFace tokenizer interface
+        tokens = self.tokenizer(seqs, padding=True, return_tensors='pt')
+        tokens = extract_input_ids(tokens, self.used_device)
+        output = self.AbLang(tokens)
+        if hasattr(output, 'last_hidden_state'):
+            return output.last_hidden_state.detach()
+        return output.detach()
+
+    def _predict_logits_with_step_masking(self, seqs):
+        # Override the stepwise masking method to use HuggingFace tokenizer
+        tokens = self.tokenizer(seqs, padding=True, return_tensors='pt')
+        tokens = extract_input_ids(tokens, self.used_device)
+        
+        logits = []
+        for single_seq_tokens in tokens:
+            tkn_len = len(single_seq_tokens)
+            masked_tokens = single_seq_tokens.repeat(tkn_len, 1)
+            for num in range(tkn_len):
+                masked_tokens[num, num] = self.tokenizer.mask_token_id
+            
+            with torch.no_grad():
+                logits_tmp = self.AbLang(masked_tokens)
+                       
+            logits_tmp = torch.stack([logits_tmp[num, num] for num in range(tkn_len)])
+            logits.append(logits_tmp)
+    
+        return torch.stack(logits, dim=0)
+
+    def _preprocess_labels(self, labels):
+        labels = extract_input_ids(labels, self.used_device)
+        return labels
+
+    def __call__(self, seqs, mode='seqcoding', align=False, stepwise_masking=False, fragmented=False, batch_size=50):
+        """
+        Use different modes for different usecases, mimicking the original pretrained class.
+        """
+        # Local implementation of format_seq_input
+        def format_seq_input(seqs, fragmented=False):
+            """Format input sequences for processing."""
+            if isinstance(seqs[0], str):
+                seqs = [seqs]
+            
+            if fragmented:
+                # For fragmented sequences, format as VH|VL without angle brackets
+                formatted_seqs = []
+                for seq in seqs:
+                    if isinstance(seq, (list, tuple)) and len(seq) == 2:
+                        heavy, light = seq[0], seq[1]
+                        formatted_seqs.append(f"{heavy}|{light}")
+                    else:
+                        formatted_seqs.append(seq)
+                return formatted_seqs, 'HL'
+            else:
+                # For non-fragmented sequences, add angle brackets: <VH>|<VL>
+                formatted_seqs = []
+                for seq in seqs:
+                    if isinstance(seq, (list, tuple)) and len(seq) == 2:
+                        heavy, light = seq[0], seq[1]
+                        # Add angle brackets and handle empty sequences
+                        heavy_part = f"<{heavy}>" if heavy else "<>"
+                        light_part = f"<{light}>" if light else "<>"
+                        formatted_seqs.append(f"{heavy_part}|{light_part}".replace("<>", ""))
+                    else:
+                        formatted_seqs.append(seq)
+                
+                return formatted_seqs, 'HL'
+
+        valid_modes = [
+            'rescoding', 'seqcoding', 'restore', 'likelihood', 'probability',
+            'pseudo_log_likelihood', 'confidence'
+        ]
+        if mode not in valid_modes:
+            raise SyntaxError(f"Given mode doesn't exist. Please select one of the following: {valid_modes}.")
+
+        seqs, chain = format_seq_input(seqs, fragmented=fragmented)
+
+        if align:
+            numbered_seqs, seqs, number_alignment = self.number_sequences(
+                seqs, chain=chain, fragmented=fragmented
+            )
+        else:
+            numbered_seqs = None
+            number_alignment = None
+
+        subset_list = []
+        for subset in [seqs[x:x+batch_size] for x in range(0, len(seqs), batch_size)]:
+            subset_list.append(getattr(self, mode)(subset, align=align, stepwise_masking=stepwise_masking))
+
+        return self.reformat_subsets(
+            subset_list,
+            mode=mode,
+            align=align,
+            numbered_seqs=numbered_seqs,
+            seqs=seqs,
+            number_alignment=number_alignment,
+        )
+
+    def pseudo_log_likelihood(self, seqs, **kwargs):
+        """
+        Original (non-vectorized) pseudo log-likelihood computation matching notebook behavior.
+        """
+        # Format input: join VH and VL with '|'
+        formatted_seqs = []
+        for s in seqs:
+            if isinstance(s, (list, tuple)):
+                formatted_seqs.append('|'.join(s))
+            else:
+                formatted_seqs.append(s)
+
+        # Tokenize all sequences in batch
+        labels = self.tokenizer(
+            formatted_seqs, padding=True, return_tensors='pt'
+        )
+        labels = extract_input_ids(labels, self.used_device)
+
+        # Convert special tokens to IDs
+        if isinstance(self.tokenizer.all_special_tokens[0], int):
+            special_token_ids = set(self.tokenizer.all_special_tokens)
+        else:
+            special_token_ids = set(self.tokenizer.convert_tokens_to_ids(tok) for tok in self.tokenizer.all_special_tokens)
+        pad_token_id = self.tokenizer.pad_token_id
+
+        mask_token_id = getattr(self.tokenizer, 'mask_token_id', None)
+        if mask_token_id is None:
+            mask_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)
+
+        plls = []
+        with torch.no_grad():
+            for i, seq_label in enumerate(labels):
+                seq_pll = []
+                for j, token_id in enumerate(seq_label):
+                    if token_id.item() in special_token_ids or token_id.item() == pad_token_id:
+                        continue
+                    masked = seq_label.clone()
+                    masked[j] = mask_token_id
+                    logits = self.AbLang(masked.unsqueeze(0))
+                    if hasattr(logits, 'last_hidden_state'):
+                        logits = logits.last_hidden_state
+                    logits = logits[0, j]
+                    nll = torch.nn.functional.cross_entropy(
+                        logits.unsqueeze(0), token_id.unsqueeze(0), reduction="none"
+                    )
+                    seq_pll.append(-nll.item())
+                if seq_pll:
+                    plls.append(np.mean(seq_pll))
+                else:
+                    plls.append(float('nan'))
+        return np.array(plls)
+
+    def seqcoding(self, seqs, **kwargs):
+        """Sequence specific representations - returns 480-dimensional embeddings for each sequence."""
+        # Format input: join VH and VL with '|'
+        formatted_seqs = []
+        for s in seqs:
+            if isinstance(s, (list, tuple)):
+                formatted_seqs.append('|'.join(s))
+            else:
+                formatted_seqs.append(s)
+        
+        # Get embeddings using the model
+        embeddings = self._encode_sequences(formatted_seqs)
+        
+        # Return sequence-level embeddings (mean pooling over sequence length)
+        # Remove batch dimension and take mean over sequence dimension
+        if len(embeddings.shape) == 3:  # [batch_size, seq_len, hidden_size]
+            seq_embeddings = embeddings.mean(dim=1)  # [batch_size, hidden_size]
+        else:
+            seq_embeddings = embeddings
+        
+        return seq_embeddings.cpu().numpy()
+
+    def rescoding(self, seqs, align=False, **kwargs):
+        """Residue specific representations - returns 480-dimensional embeddings for each residue."""
+        # Format input: join VH and VL with '|'
+        formatted_seqs = []
+        for s in seqs:
+            if isinstance(s, (list, tuple)):
+                formatted_seqs.append('|'.join(s))
+            else:
+                formatted_seqs.append(s)
+        
+        # Get embeddings using the model
+        embeddings = self._encode_sequences(formatted_seqs)
+        
+        # Return residue-level embeddings
+        # embeddings shape: [batch_size, seq_len, hidden_size]
+        if len(embeddings.shape) == 3:
+            # Convert to numpy and return as list of arrays for each sequence
+            embeddings_np = embeddings.cpu().numpy()
+            return [embeddings_np[i] for i in range(embeddings_np.shape[0])]
+        else:
+            return embeddings.cpu().numpy()
+
+    def likelihood(self, seqs, align=False, stepwise_masking=False, **kwargs):
+        """Likelihood of mutations - returns logits for each amino acid at each position."""
+        # Format input: join VH and VL with '|'
+        formatted_seqs = []
+        for s in seqs:
+            if isinstance(s, (list, tuple)):
+                formatted_seqs.append('|'.join(s))
+            else:
+                formatted_seqs.append(s)
+
+        # Get logits
+        if stepwise_masking:
+            logits = self._predict_logits_with_step_masking(formatted_seqs)
+        else:
+            logits = self._predict_logits(formatted_seqs)
+        
+        # Return logits as numpy array
+        return logits.cpu().numpy()
+
+    def confidence(self, seqs, **kwargs):
+        """Confidence calculation - match original ablang2 implementation by excluding all special tokens from loss."""
+        # Format input: join VH and VL with '|'
+        formatted_seqs = []
+        for s in seqs:
+            if isinstance(s, (list, tuple)):
+                formatted_seqs.append('|'.join(s))
+            else:
+                formatted_seqs.append(s)
+        
+        plls = []
+        for seq in formatted_seqs:
+            tokens = self.tokenizer([seq], padding=True, return_tensors='pt')
+            input_ids = extract_input_ids(tokens, self.used_device)
+            
+            with torch.no_grad():
+                output = self.AbLang(input_ids)
+                if hasattr(output, 'last_hidden_state'):
+                    logits = output.last_hidden_state
+                else:
+                    logits = output
+                
+                # Get the sequence (remove batch dimension)
+                logits = logits[0]  # [seq_len, vocab_size]
+                input_ids = input_ids[0]  # [seq_len]
+                
+                # Exclude all special tokens (pad, mask, etc.)
+                if isinstance(self.tokenizer.all_special_tokens[0], int):
+                    special_token_ids = set(self.tokenizer.all_special_tokens)
+                else:
+                    special_token_ids = set(self.tokenizer.convert_tokens_to_ids(tok) for tok in self.tokenizer.all_special_tokens)
+                valid_mask = ~torch.isin(input_ids, torch.tensor(list(special_token_ids), device=input_ids.device))
+                
+                if valid_mask.sum() > 0:
+                    valid_logits = logits[valid_mask]
+                    valid_labels = input_ids[valid_mask]
+                    
+                    # Calculate cross-entropy loss
+                    nll = torch.nn.functional.cross_entropy(
+                        valid_logits,
+                        valid_labels,
+                        reduction="mean"
+                    )
+                    pll = -nll.item()
+                else:
+                    pll = 0.0
+                
+                plls.append(pll)
+        
+        return np.array(plls, dtype=np.float32)
+
+    def probability(self, seqs, align=False, stepwise_masking=False, **kwargs):
+        """
+        Probability of mutations - applies softmax to logits to get probabilities
+        """
+        # Format input: join VH and VL with '|'
+        formatted_seqs = []
+        for s in seqs:
+            if isinstance(s, (list, tuple)):
+                formatted_seqs.append('|'.join(s))
+            else:
+                formatted_seqs.append(s)
+
+        # Get logits
+        if stepwise_masking:
+            # For stepwise masking, we need to implement it similar to likelihood
+            # This is a simplified version - you might want to implement full stepwise masking
+            logits = self._predict_logits(formatted_seqs)
+        else:
+            logits = self._predict_logits(formatted_seqs)
+        
+        # Apply softmax to get probabilities
+        probs = logits.softmax(-1).cpu().numpy()
+        
+        if align:
+            return probs
+        else:
+            # Return residue-level probabilities (excluding special tokens)
+            return [res_to_list(state, seq) for state, seq in zip(probs, formatted_seqs)]
+
+    def restore(self, seqs, align=False, **kwargs):
+        hf_abrestore = HFAbRestore(self.AbLang, self.tokenizer, spread=self.spread, device=self.used_device, ncpu=self.ncpu)
+        restored = hf_abrestore.restore(seqs, align=align)
+        # Apply angle brackets formatting to match original format
+        if isinstance(restored, np.ndarray):
+            restored = np.array([add_angle_brackets(seq) for seq in restored])
+        else:
+            restored = [add_angle_brackets(seq) for seq in restored]
+        return restored
+
+def extract_input_ids(tokens, device):
+    if hasattr(tokens, 'input_ids'):
+        return tokens.input_ids.to(device)
+    elif isinstance(tokens, dict):
+        if 'input_ids' in tokens:
+            return tokens['input_ids'].to(device)
+        else:
+            for v in tokens.values():
+                if hasattr(v, 'ndim') or torch.is_tensor(v):
+                    return v.to(device)
+    elif torch.is_tensor(tokens):
+        return tokens.to(device)
+    else:
+        raise ValueError("Could not extract input_ids from tokenizer output")

alignment.py

@@ -0,0 +1,86 @@
+from dataclasses import dataclass
+import numpy as np
+import torch
+
+from extra_utils import paired_msa_numbering, unpaired_msa_numbering, create_alignment
+
+
+class AbAlignment:
+
+    def __init__(self, device = 'cpu', ncpu = 1):
+        
+        self.device = device
+        self.ncpu = ncpu
+        
+    def number_sequences(self, seqs, chain = 'H', fragmented = False):
+        if chain == 'HL':
+            numbered_seqs, seqs, number_alignment = paired_msa_numbering(seqs, fragmented = fragmented, n_jobs = self.ncpu)
+        else:
+            assert chain == 'HL', 'Currently "Align==True" only works for paired sequences. \nPlease use paired sequences or Align=False.'
+            numbered_seqs, seqs, number_alignment = unpaired_msa_numbering(
+                seqs, chain = chain, fragmented = fragmented, n_jobs = self.ncpu
+            )
+        
+        return numbered_seqs, seqs, number_alignment
+    
+    def align_encodings(self, encodings, numbered_seqs, seqs, number_alignment):
+        
+        aligned_list = [
+            create_alignment(
+                res_embed, numbered_seq, seq, number_alignment
+            ) for res_embed, numbered_seq, seq in zip(encodings, numbered_seqs, seqs)
+        ]
+        aligned_encodings = np.concatenate([aligned_list], axis=0)
+        return aligned_encodings
+        
+        
+    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'
+        ]:
+            return np.concatenate(subset_list)
+        elif align:
+            subset_list = [
+                self.align_encodings(
+                    subset, 
+                    numbered_seqs[num*len(subset):(num+1)*len(subset)],
+                    seqs[num*len(subset):(num+1)*len(subset)], 
+                    number_alignment
+                ) for num, subset in enumerate(subset_list)
+            ]       
+            
+            subset = np.concatenate(subset_list)
+            
+            return aligned_results(
+                aligned_seqs = [''.join(alist) for alist in subset[:,:,-1]],
+                aligned_embeds = subset[:,:,:-1].astype(float),
+                number_alignment=number_alignment.apply(lambda x: '{}{}'.format(*x[0]), axis=1).values
+            ) 
+    
+        elif not align:
+            return sum(subset_list, [])
+        else:
+            return np.concatenate(subset_list) # this needs to be changed
+        
+
+@dataclass
+class aligned_results():
+    """
+    Dataclass used to store output.
+    """
+    
+    aligned_seqs: None
+    aligned_embeds: None
+    number_alignment: None

app.py

@@ -0,0 +1,336 @@
+import gradio as gr
+import sys
+import os
+from transformers import AutoModel, AutoTokenizer
+from transformers.utils import cached_file
+
+# Load model and tokenizer from Hugging Face Hub
+model = AutoModel.from_pretrained("hemantn/ablang2", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("hemantn/ablang2", trust_remote_code=True)
+
+# Find the cached model directory and import adapter
+adapter_path = cached_file("hemantn/ablang2", "adapter.py")
+cached_model_dir = os.path.dirname(adapter_path)
+sys.path.insert(0, cached_model_dir)
+
+# Import and create the adapter
+from adapter import AbLang2PairedHuggingFaceAdapter
+ablang = AbLang2PairedHuggingFaceAdapter(model=model, tokenizer=tokenizer)
+
+def restore_sequences(heavy_chain, light_chain, use_align=False):
+    """
+    Restore masked residues in antibody sequences.
+    
+    Args:
+        heavy_chain (str): Heavy chain sequence with masked residues (*)
+        light_chain (str): Light chain sequence with masked residues (*)
+        use_align (bool): Whether to use alignment for variable missing lengths
+    
+    Returns:
+        tuple: (restored_heavy, restored_light, highlighted_heavy, highlighted_light)
+    """
+    try:
+        # Check if alignment is requested but not available
+        if use_align:
+            try:
+                import anarci
+            except ImportError:
+                return "Alignment feature requires 'anarci' package which is not available. Please disable alignment option.", "", ""
+        # Prepare input sequences
+        if heavy_chain.strip() and light_chain.strip():
+            # Both chains provided
+            sequences = [[heavy_chain.strip(), light_chain.strip()]]
+        elif heavy_chain.strip():
+            # Only heavy chain provided
+            sequences = [[heavy_chain.strip(), ""]]
+        elif light_chain.strip():
+            # Only light chain provided
+            sequences = [["", light_chain.strip()]]
+        else:
+            return "Please provide at least one antibody chain sequence.", "", "", ""
+        
+        # Perform restoration
+        restored = ablang(sequences, mode='restore', align=use_align)
+        
+        # Format output
+        if hasattr(restored, '__len__') and len(restored) > 0:
+            result = restored[0]  # Get the first (and only) result
+            
+            # Parse the result to separate heavy and light chains
+            if '>|<' in result:
+                # Both chains present
+                heavy_part = result.split('>|<')[0].replace('<', '').replace('>', '')
+                light_part = result.split('>|<')[1].replace('<', '').replace('>', '')
+            elif result.startswith('<') and result.endswith('>'):
+                # Only one chain present
+                if heavy_chain.strip():
+                    heavy_part = result.replace('<', '').replace('>', '')
+                    light_part = ""
+                else:
+                    heavy_part = ""
+                    light_part = result.replace('<', '').replace('>', '')
+            else:
+                return "Error: Unexpected result format.", "", "", ""
+            
+            # Create highlighted versions
+            highlighted_heavy = highlight_restored_residues(heavy_chain.strip(), heavy_part)
+            highlighted_light = highlight_restored_residues(light_chain.strip(), light_part)
+            
+            # Create HTML outputs with proper styling - no scroll, wrap text
+            heavy_html = f'<div class="restored-sequence-box" style="padding: 10px; background-color: #f8f9fa; border: 1px solid #dee2e6; border-radius: 4px;">{highlighted_heavy}</div>'
+            light_html = f'<div class="restored-sequence-box" style="padding: 10px; background-color: #f8f9fa; border: 1px solid #dee2e6; border-radius: 4px;">{highlighted_light}</div>'
+            
+            return heavy_html, light_html
+        else:
+            return "Error: No restoration result obtained.", "", ""
+            
+    except Exception as e:
+        return f"Error during restoration: {str(e)}", "", ""
+
+def highlight_restored_residues(original_seq, restored_seq):
+    """
+    Highlight restored residues in green.
+    """
+    if not original_seq or not restored_seq:
+        return restored_seq
+    
+    highlighted = ""
+    for i, (orig_char, rest_char) in enumerate(zip(original_seq, restored_seq)):
+        if orig_char == '*' and rest_char != '*':
+            # This residue was restored
+            highlighted += f'<span class="restored-highlight">{rest_char}</span>'
+        else:
+            highlighted += rest_char
+    
+    # Add any remaining characters from restored sequence
+    if len(restored_seq) > len(original_seq):
+        highlighted += restored_seq[len(original_seq):]
+    
+    return highlighted
+
+# Create Gradio interface
+with gr.Blocks(title="AbLang2 Sequence Restorer", theme=gr.themes.Soft(), css="""
+    * {
+        font-family: 'Courier New', monospace !important;
+    }
+    .sequence-input, .sequence-output {
+        font-family: 'Courier New', monospace !important;
+        font-size: 14px !important;
+        letter-spacing: 0.5px !important;
+    }
+    .restored-highlight {
+        background-color: #90EE90 !important;
+        color: #000 !important;
+        font-weight: bold !important;
+    }
+    .examples {
+        font-family: 'Courier New', monospace !important;
+        font-size: 14px !important;
+        letter-spacing: 0.5px !important;
+    }
+    .restored-sequence-box {
+        font-family: 'Courier New', monospace !important;
+        font-size: 14px !important;
+        letter-spacing: 0.5px !important;
+        white-space: pre-wrap !important;
+        word-wrap: break-word !important;
+        overflow-wrap: break-word !important;
+    }
+    .restored-heading {
+        color: #2E8B57 !important;
+        font-weight: bold !important;
+        font-size: 18px !important;
+    }
+    .example-text {
+        font-family: 'Courier New', monospace !important;
+        font-size: 12px !important;
+        white-space: pre-wrap !important;
+        word-wrap: break-word !important;
+    }
+    .examples-table {
+        font-family: 'Courier New', monospace !important;
+        font-size: 12px !important;
+        white-space: pre-wrap !important;
+        word-wrap: break-word !important;
+        max-width: none !important;
+        overflow: visible !important;
+    }
+    .examples-table td {
+        font-family: 'Courier New', monospace !important;
+        font-size: 12px !important;
+        white-space: pre-wrap !important;
+        word-wrap: break-word !important;
+        max-width: none !important;
+        overflow: visible !important;
+        text-overflow: unset !important;
+    }
+    .sequence-output label {
+        font-weight: bold !important;
+        color: #495057 !important;
+        font-size: 14px !important;
+        margin-bottom: 5px !important;
+    }
+    /* Force full display of examples */
+    .examples-container {
+        font-family: 'Courier New', monospace !important;
+        font-size: 12px !important;
+    }
+    .examples-container table {
+        width: 100% !important;
+        table-layout: auto !important;
+    }
+    .examples-container td {
+        white-space: pre-wrap !important;
+        word-wrap: break-word !important;
+        overflow-wrap: break-word !important;
+        max-width: none !important;
+        text-overflow: unset !important;
+        padding: 8px !important;
+        vertical-align: top !important;
+    }
+    .examples-container th {
+        white-space: nowrap !important;
+        padding: 8px !important;
+    }
+    /* Override any Gradio default truncation */
+    .examples table td {
+        white-space: pre-wrap !important;
+        word-wrap: break-word !important;
+        overflow-wrap: break-word !important;
+        max-width: none !important;
+        text-overflow: unset !important;
+        overflow: visible !important;
+        font-family: 'Courier New', monospace !important;
+        font-size: 12px !important;
+    }
+    .examples table {
+        table-layout: auto !important;
+        width: 100% !important;
+    }
+    /* Target the specific examples component */
+    div[data-testid="examples"] table td {
+        white-space: pre-wrap !important;
+        word-wrap: break-word !important;
+        overflow-wrap: break-word !important;
+        max-width: none !important;
+        text-overflow: unset !important;
+        overflow: visible !important;
+        font-family: 'Courier New', monospace !important;
+        font-size: 12px !important;
+    }
+    /* Force examples to show full content */
+    .examples table, .examples table td, .examples table th {
+        white-space: pre-wrap !important;
+        word-wrap: break-word !important;
+        overflow-wrap: break-word !important;
+        max-width: none !important;
+        text-overflow: unset !important;
+        overflow: visible !important;
+        font-family: 'Courier New', monospace !important;
+        font-size: 12px !important;
+        table-layout: auto !important;
+        width: auto !important;
+        min-width: 100% !important;
+    }
+    /* Override any inline styles */
+    .examples * {
+        white-space: pre-wrap !important;
+        word-wrap: break-word !important;
+        overflow-wrap: break-word !important;
+        max-width: none !important;
+        text-overflow: unset !important;
+        overflow: visible !important;
+    }
+    /* Style output labels to match input labels exactly */
+    .output-label {
+        font-weight: 600 !important;
+        color: var(--label-text-color) !important;
+        font-size: 14px !important;
+        margin-bottom: 8px !important;
+        margin-top: 16px !important;
+        line-height: 1.4 !important;
+        display: block !important;
+    }
+""") as demo:
+    gr.Markdown("""
+    # 🧬 AbLang2 Sequence Restorer
+    
+    This app uses the AbLang2 model to restore masked residues (*) in antibody sequences. 
+    You can provide either one or both heavy and light chain sequences.
+    
+    **Instructions:**
+    - Use `*` to mask residues you want to restore
+    - Provide heavy chain, light chain, or both
+    - Enable "Use Alignment" for variable missing lengths
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            heavy_input = gr.Textbox(
+                label="Heavy Chain Sequence",
+                placeholder="Enter heavy chain sequence with masked residues (*)...",
+                lines=3,
+                max_lines=5,
+                elem_classes=["sequence-input"]
+            )
+            
+            light_input = gr.Textbox(
+                label="Light Chain Sequence", 
+                placeholder="Enter light chain sequence with masked residues (*)...",
+                lines=3,
+                max_lines=5,
+                elem_classes=["sequence-input"]
+            )
+            
+            align_checkbox = gr.Checkbox(
+                label="Use Alignment (for variable missing lengths) - Requires anarci package",
+                value=False
+            )
+            
+            restore_btn = gr.Button("🔄 Restore Sequences", variant="primary")
+        
+        with gr.Column():
+            gr.Markdown("### 🧬 Restored Sequences", elem_classes=["restored-heading"])
+            gr.Markdown("*Green highlighting shows restored residues*")
+            
+            gr.Markdown("**Heavy Chain Sequence**", elem_classes=["output-label"])
+            heavy_output = gr.HTML(label="")
+            
+            gr.Markdown("**Light Chain Sequence**", elem_classes=["output-label"])
+            light_output = gr.HTML(label="")
+    
+    # Example sequences
+    gr.Examples(
+        examples=[
+            [
+                "EVQ***SGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCAR**PGHGAAFMDVWGTGTTVTVSS",
+                "DIQLTQSPLSLPVTLGQPASISCRSS*SLEASDTNIYLSWFQQRPGQSPRRLIYKI*NRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK"
+            ],
+            [
+                "EVQLVESGGGLVQPGGSLRLSCAASGFTFSSYAMGWVRQAPGKGLEWVSAISGSGGSTYYADSVKGRFTISRDNSKNTLYLQMNSLRAEDTAVYYCARDY**GMDVWGQGTTVTVSS",
+                ""
+            ],
+            [
+                "",
+                "DIQLTQSPSSLSASVGDRVTITCRASQSISSYLNWYQQKPGKAPKLLIY*ASSLQSGVPSRFSGSGSGTDFTLTISSLQPEDFATYYCQQSYSTP*TFGQGTKVEIK"
+            ]
+        ],
+        inputs=[heavy_input, light_input],
+        label="Example Sequences"
+    )
+    
+    # Connect the button to the function
+    restore_btn.click(
+        fn=restore_sequences,
+        inputs=[heavy_input, light_input, align_checkbox],
+        outputs=[heavy_output, light_output]
+    )
+    
+    gr.Markdown("""
+    ---
+    **Note:** This app uses the AbLang2 model from Hugging Face Hub. 
+    The restoration process may take a few seconds depending on sequence length and complexity.
+    """)
+
+if __name__ == "__main__":
+    demo.launch() 

config.json

@@ -0,0 +1,23 @@
+{
+  "model_type": "ablang2-paired",
+  "vocab_size": 26,
+  "hidden_embed_size": 480,
+  "n_attn_heads": 20,
+  "n_encoder_blocks": 12,
+  "padding_tkn": 21,
+  "mask_tkn": 23,
+  "layer_norm_eps": 1e-12,
+  "a_fn": "swiglu",
+  "dropout": 0.0,
+  "tokenizer_class": "AbLang2PairedTokenizer",
+  "auto_map": {
+    "AutoConfig": "configuration_ablang2paired.AbLang2PairedConfig",
+    "AutoModel": "modeling_ablang2paired.AbLang2PairedHFModel",
+    "AutoTokenizer": ["tokenizer_ablang2paired.AbLang2PairedTokenizer", "tokenizer_ablang2paired.AbLang2PairedTokenizer"],
+    "AbRestore": "restoration.AbRestore",
+    "AbScores": "scores.AbScores",
+    "AbAlignment": "alignment.AbAlignment",
+    "AbEncoding": "ablang_encodings.AbEncoding",
+    "ExtraUtils": "extra_utils"
+  }
+}

configuration_ablang2paired.py

@@ -0,0 +1,31 @@
+from transformers import PretrainedConfig
+
+class AbLang2PairedConfig(PretrainedConfig):
+    model_type = "ablang2-paired"
+
+    def __init__(
+        self,
+        vocab_size=26,
+        hidden_embed_size=480,
+        n_attn_heads=20,
+        n_encoder_blocks=12,
+        padding_tkn=21,
+        mask_tkn=23,
+        layer_norm_eps=1e-12,
+        a_fn="swiglu",
+        dropout=0.0,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.hidden_embed_size = hidden_embed_size
+        self.hidden_size = hidden_embed_size  # Add this for Hugging Face compatibility
+        self.n_attn_heads = n_attn_heads
+        self.num_attention_heads = n_attn_heads  # Add this for Hugging Face compatibility
+        self.num_hidden_layers = n_encoder_blocks  # Add this for Hugging Face compatibility
+        self.n_encoder_blocks = n_encoder_blocks
+        self.padding_tkn = padding_tkn
+        self.mask_tkn = mask_tkn
+        self.layer_norm_eps = layer_norm_eps
+        self.a_fn = a_fn
+        self.dropout = dropout 

encoderblock.py

@@ -0,0 +1,173 @@
+import torch
+import math
+from torch import nn
+import torch.nn.functional as F
+import einops
+from rotary_embedding_torch import RotaryEmbedding
+
+class TransformerEncoder(torch.nn.Module):
+    """
+    Single Transformer Encoder.
+    
+    """
+    def __init__(
+        self, 
+        hidden_embed_size,
+        n_attn_heads,
+        attn_dropout: float = 0.0,
+        layer_norm_eps: float = 1e-05,
+        a_fn: str = "gelu",
+    ):
+        super().__init__()
+        
+        assert hidden_embed_size % n_attn_heads == 0, \
+        "Embedding dimension must be devisible with the number of heads." 
+        
+        self.multihead_attention = MultiHeadAttention(
+            embed_dim = hidden_embed_size, 
+            num_heads = n_attn_heads,
+            attention_dropout_prob = attn_dropout
+        )
+        
+        activation_fn, scale = get_activation_fn(a_fn)
+        
+        self.intermediate_layer = torch.nn.Sequential(
+            torch.nn.Linear(hidden_embed_size, hidden_embed_size * 4 * scale),
+            activation_fn(),
+            torch.nn.Linear(hidden_embed_size * 4, hidden_embed_size),
+        )
+        
+        self.pre_attn_layer_norm = torch.nn.LayerNorm(hidden_embed_size, eps=layer_norm_eps)
+        self.final_layer_norm = torch.nn.LayerNorm(hidden_embed_size, eps=layer_norm_eps)
+        
+    def forward(self, hidden_embed, attn_mask=None, return_attn_weights: bool = False):
+        
+        residual = hidden_embed
+        hidden_embed = self.pre_attn_layer_norm(hidden_embed.clone())
+        hidden_embed, attn_weights = self.multihead_attention(
+            hidden_embed, 
+            attn_mask=attn_mask, 
+            return_attn_weights=return_attn_weights
+        )
+        hidden_embed = residual + hidden_embed
+        
+        residual = hidden_embed
+        hidden_embed = self.final_layer_norm(hidden_embed)
+        hidden_embed = self.intermediate_layer(hidden_embed)
+        hidden_embed = residual + hidden_embed
+        return hidden_embed, attn_weights   
+    
+class MultiHeadAttention(torch.nn.Module):
+
+    def __init__(
+        self, 
+        embed_dim,
+        num_heads, 
+        attention_dropout_prob: float = 0.0,
+        bias: bool = True,
+    ):
+        super().__init__()
+        
+        self.attention_dropout = torch.nn.Dropout(attention_dropout_prob)
+
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        assert (self.head_dim * num_heads == self.embed_dim), "embed_dim must be divisible by num_heads"
+        self.scaling = self.head_dim**-0.5
+        
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+            
+        self.reset_parameters()
+        
+        self.rotary_emb = RotaryEmbedding(dim = self.head_dim)
+        
+    def reset_parameters(self):
+        
+        nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))
+        nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))
+        nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2))
+
+        nn.init.xavier_uniform_(self.out_proj.weight)
+        if self.out_proj.bias is not None:
+            nn.init.constant_(self.out_proj.bias, 0.0)
+        
+    def attention(self, q, k, v, attn_mask=None):
+        
+        attn_weights = torch.matmul(q, k.transpose(-2, -1)) 
+        attn_weights = attn_weights / math.sqrt(self.head_dim)
+                
+        if attn_mask is not None:
+            attn_mask = einops.rearrange(
+                attn_mask, 
+                'b_size (h1 h2 seq_len) -> b_size h1 h2 seq_len', 
+                h1=1, h2=1
+            )
+            attn_weights = attn_weights.masked_fill(attn_mask, float("-inf"))
+
+        attn_weights = F.softmax(attn_weights, dim=-1)
+        
+        attn = self.attention_dropout(attn_weights)
+        attn = torch.matmul(attn, v)
+        return attn, attn_weights       
+
+    def forward(self, x, attn_mask=None, return_attn_weights: bool = False):
+        
+        batch_size, seq_len, embed_dim = x.size()
+        
+        q, k, v = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+        q *= self.scaling
+        
+        q = q.contiguous().view(
+            batch_size, 
+            seq_len, 
+            self.num_heads, 
+            self.head_dim
+        ).transpose(1, 2) # [n_batch, n_heads, seq_len, head_dim]
+        k = k.contiguous().view(
+            batch_size, 
+            seq_len, 
+            self.num_heads, 
+            self.head_dim
+        ).transpose(1, 2) # [n_batch, n_heads, seq_len, head_dim]
+        v = v.contiguous().view(
+            batch_size, 
+            seq_len, 
+            self.num_heads, 
+            self.head_dim
+        ).transpose(1, 2) # [n_batch, n_heads, seq_len, head_dim]
+        
+        q = self.rotary_emb.rotate_queries_or_keys(q)
+        k = self.rotary_emb.rotate_queries_or_keys(k)
+        
+        # Determine value outputs
+        attn, attn_weights = self.attention(
+            q, k, v, 
+            attn_mask=attn_mask
+        ) # attn_weights [n_batch, n_heads, seq_len (target), seq_len (source)]
+    
+        attn = attn.transpose(1, 2).reshape(batch_size, seq_len, embed_dim)       
+        attn = self.out_proj(attn)
+
+        if return_attn_weights:
+            return attn, attn_weights
+        else:
+            return attn, None
+        
+class SwiGLU(torch.nn.Module):
+    def forward(self, x):
+        x, gate = x.chunk(2, dim=-1) 
+        return F.silu(gate) * x
+        
+def get_activation_fn(a_fn):
+    
+    if a_fn == "gelu":
+        return torch.nn.GELU, 1
+    
+    elif a_fn == "swiglu":
+        return SwiGLU, 2
+    

environment.yaml

@@ -0,0 +1,44 @@
+name: AbLang
+channels:
+  - conda-forge
+  - pytorch
+  - bioconda
+  - defaults
+dependencies:
+  - python=3.10.18
+  - pip
+  - pytorch=2.5.1
+  - pytorch-cuda=12.4
+  - numpy=2.2.6
+  - pandas=2.3.1
+  - transformers=4.53.3
+  - anarci=2024.05.21
+  - jupyter=7.4.4
+  - notebook=7.4.4
+  - ipython=8.37.0
+  - ipykernel=6.29.5
+  - matplotlib-inline=0.1.7
+  - scikit-learn
+  - matplotlib
+  - seaborn
+  - biopython=1.85
+  - huggingface_hub=0.33.4
+  - tokenizers=0.21.3
+  - safetensors=0.5.3
+  - einops=0.8.1
+  - tqdm=4.67.1
+  - requests=2.32.4
+  - urllib3=2.5.0
+  - certifi=2025.7.14
+  - filelock=3.18.0
+  - fsspec=2025.3.0
+  - packaging=25.0
+  - regex=2024.11.6
+  - sympy=1.13.3
+  - networkx=3.4.2
+  - jinja2=3.1.6
+  - pyyaml=6.0.2
+  - typing_extensions=4.14.1
+  - pip:
+    - numba=0.61.2
+    - llvmlite=0.44.0 

extra_utils.py

@@ -0,0 +1,165 @@
+import string, re
+import numpy as np
+
+
+def res_to_list(logits, seq):
+    return logits[:len(seq)]
+
+def res_to_seq(a, mode='mean'):
+    """
+    Function for how we go from n_values for each amino acid to n_values for each sequence.
+    
+    We leave out padding tokens.
+    """
+    
+    if mode=='sum':
+        return a[0:(int(a[-1]))].sum()
+    
+    elif mode=='mean':
+        return a[0:(int(a[-1]))].mean()
+    
+    elif mode=='restore':
+        return a[0][0:(int(a[-1]))]
+
+def get_number_alignment(numbered_seqs):
+    """
+    Creates a number alignment from the anarci results.
+    """
+    import pandas as pd
+
+    alist = [pd.DataFrame(aligned_seq, columns = [0,1,'resi']) for aligned_seq in numbered_seqs]
+    unsorted_alignment = pd.concat(alist).drop_duplicates(subset=0)
+    max_alignment = get_max_alignment()    
+    
+    return max_alignment.merge(unsorted_alignment.query("resi!='-'"), left_on=0, right_on=0)[[0,1]]
+
+def get_max_alignment():
+    """
+    Create maximum possible alignment for sorting
+    """
+    import pandas as pd
+
+    sortlist = [[("<", "")]]
+    for num in range(1, 128+1):
+        if num in [33,61,112]:
+            for char in string.ascii_uppercase[::-1]:
+                sortlist.append([(num, char)])
+
+            sortlist.append([(num,' ')])
+        else:
+            sortlist.append([(num,' ')])
+            for char in string.ascii_uppercase:
+                sortlist.append([(num, char)])
+                
+    return pd.DataFrame(sortlist + [[(">", "")]])
+
+
+def paired_msa_numbering(ab_seqs, fragmented = False, n_jobs = 10):
+    
+    import pandas as pd
+    
+    tmp_seqs = [pairs.replace(">", "").replace("<", "").split("|") for pairs in ab_seqs]
+
+    numbered_seqs_heavy, seqs_heavy, number_alignment_heavy = unpaired_msa_numbering(
+        [i[0] for i in tmp_seqs], 'H', fragmented = fragmented, n_jobs = n_jobs
+    )
+    numbered_seqs_light, seqs_light, number_alignment_light = unpaired_msa_numbering(
+        [i[1] for i in tmp_seqs], 'L', fragmented = fragmented, n_jobs = n_jobs
+    )
+    
+    number_alignment = pd.concat([
+        number_alignment_heavy, 
+        pd.DataFrame([[("|",""), "|"]]), 
+        number_alignment_light]
+    ).reset_index(drop=True)
+    
+    seqs = [f"{heavy}|{light}" for heavy, light in zip(seqs_heavy, seqs_light)]
+    numbered_seqs = [
+        heavy + [(("|",""), "|", "|")] + light for heavy, light in zip(numbered_seqs_heavy, numbered_seqs_light)
+    ]
+    
+    return numbered_seqs, seqs, number_alignment
+
+
+def unpaired_msa_numbering(seqs, chain = 'H', fragmented = False, n_jobs = 10):
+    
+    numbered_seqs = number_with_anarci(seqs, chain = chain, fragmented = fragmented, n_jobs = n_jobs)
+    number_alignment = get_number_alignment(numbered_seqs)
+    number_alignment[1] = chain   
+    
+    seqs = [''.join([i[2] for i in numbered_seq]).replace('-','') for numbered_seq in numbered_seqs]    
+    return numbered_seqs, seqs, number_alignment
+
+
+def number_with_anarci(seqs, chain = 'H', fragmented = False, n_jobs = 1):
+    
+    import anarci
+    import pandas as pd
+
+    anarci_out = anarci.run_anarci(
+        pd.DataFrame(seqs).reset_index().values.tolist(), 
+        ncpu=n_jobs, 
+        scheme='imgt',
+        allowed_species=['human', 'mouse'],
+    )
+    
+    numbered_seqs = []
+    for onarci in anarci_out[1]:            
+        numbered_seq = []
+        for i in onarci[0][0]:
+            if i[1] != '-':
+                numbered_seq.append((i[0], chain, i[1]))
+            
+        if fragmented:
+            numbered_seqs.append(numbered_seq) 
+        else:
+            numbered_seqs.append([(("<",""), chain, "<")] + numbered_seq + [((">",""), chain, ">")])
+    
+    return numbered_seqs
+
+
+def create_alignment(res_embeds, numbered_seqs, seq, number_alignment):
+    
+    import pandas as pd
+    
+    datadf = pd.DataFrame(numbered_seqs)
+    sequence_alignment = number_alignment.merge(datadf, how='left', on=[0, 1]).fillna('-')[2]
+
+    idxs = np.where(sequence_alignment.values == '-')[0]
+    idxs = [idx-num for num, idx in enumerate(idxs)]
+    
+    aligned_embeds = pd.DataFrame(np.insert(res_embeds[:len(seq)], idxs , 0, axis=0))
+    
+    return pd.concat([aligned_embeds, sequence_alignment], axis=1).values
+
+
+def get_spread_sequences(seq, spread, start_position):
+    """
+    Test sequences which are 8 positions shorter (position 10 + max CDR1 gap of 7) up to 2 positions longer (possible insertions).
+    """
+    spread_sequences = []
+
+    for diff in range(start_position-8, start_position+2+1):
+        spread_sequences.append('*'*diff+seq)
+    
+    return np.array(spread_sequences)
+
+def get_sequences_from_anarci(out_anarci, max_position, spread):
+    """
+    Ensures correct masking on each side of sequence
+    """
+    
+    if out_anarci == 'ANARCI_error':
+        return np.array(['ANARCI-ERR']*spread)
+    
+    end_position = int(re.search(r'\d+', out_anarci[::-1]).group()[::-1])
+    # Fixes ANARCI error of poor numbering of the CDR1 region
+    start_position = int(re.search(r'\d+,\s\'.\'\),\s\'[^-]+\'\),\s\(\(\d+,\s\'.\'\),\s\'[^-]+\'\),\s\(\(\d+,\s\'.\'\),\s\'[^-]+\'\),\s\(\(\d+,\s\'.\'\),\s\'[^-]+',
+                                   out_anarci).group().split(',')[0]) - 1
+    
+    sequence = "".join(re.findall(r"(?i)[A-Z*]", "".join(re.findall(r'\),\s\'[A-Z*]', out_anarci))))
+
+    sequence_j = ''.join(sequence).replace('-','').replace('X','*') + '*'*(max_position-int(end_position))
+
+    return get_spread_sequences(sequence_j, spread, start_position)
+

hparams.json

@@ -0,0 +1 @@
+{"name": "AbLang-2", "n_encoder_blocks": 12, "hidden_embed_size": 480, "n_attn_heads": 20, "a_fn": "swiglu", "layer_norm_eps": 1e-12, "pad_tkn": 21, "start_tkn": 0, "end_tkn": 22, "sep_tkn": 25, "mask_tkn": 23, "vocab_size": 26}

model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:56d6f07862a6f824f88c8707bbc03e4026c9db762be2d3041e9767e2e6f86386
+size 179314477

modeling_ablang2paired.py

@@ -0,0 +1,116 @@
+import torch
+import os
+from torch import nn
+from transformers import PreTrainedModel
+
+# Import configuration
+try:
+    from .configuration_ablang2paired import AbLang2PairedConfig
+except ImportError:
+    from configuration_ablang2paired import AbLang2PairedConfig
+
+# Import the AbLang model from local files
+try:
+    from ablang import AbLang
+except ImportError:
+    # Fallback: try to import from the current directory
+    try:
+        from .ablang import AbLang
+    except ImportError:
+        raise ImportError(
+            "Could not find AbLang module. Please ensure ablang.py is present in the repository."
+        )
+
+
+
+
+
+class AbLang2PairedHFModel(PreTrainedModel):
+    config_class = AbLang2PairedConfig
+    model_type = "ablang2-paired"
+
+    def __init__(self, config: AbLang2PairedConfig):
+        super().__init__(config)
+        self.model = AbLang(
+            vocab_size=config.vocab_size,
+            hidden_embed_size=config.hidden_embed_size,
+            n_attn_heads=config.n_attn_heads,
+            n_encoder_blocks=config.n_encoder_blocks,
+            padding_tkn=config.padding_tkn,
+            mask_tkn=config.mask_tkn,
+            layer_norm_eps=config.layer_norm_eps,
+            a_fn=config.a_fn,
+            dropout=config.dropout,
+        )
+
+    def forward(self, input_ids=None, x=None, attention_mask=None, **kwargs):
+        # Handle both Hugging Face format (input_ids) and original format (x)
+        if input_ids is not None:
+            x = input_ids
+        elif x is None:
+            raise ValueError("Either input_ids or x must be provided")
+        
+        # Get the output from the underlying model
+        output = self.model(x, attention_mask)
+        
+        # Return as a simple object with last_hidden_state attribute
+        class ModelOutput:
+            def __init__(self, last_hidden_state):
+                self.last_hidden_state = last_hidden_state
+        
+        return ModelOutput(output)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        # Load config first
+        config = kwargs.get("config")
+        if config is None:
+            from transformers import AutoConfig
+            config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)
+        
+        # Create model with config
+        model = cls(config)
+        
+        # Try to load custom weights
+        try:
+            from transformers.utils import cached_file
+            custom_weights_path = cached_file(
+                pretrained_model_name_or_path,
+                "model.pt",
+                cache_dir=kwargs.get("cache_dir"),
+                force_download=kwargs.get("force_download", False),
+                resume_download=kwargs.get("resume_download", False),
+                proxies=kwargs.get("proxies"),
+                token=kwargs.get("token"),
+                revision=kwargs.get("revision"),
+                local_files_only=kwargs.get("local_files_only", False),
+            )
+            
+            if custom_weights_path is not None and os.path.exists(custom_weights_path):
+                # Load custom weights
+                state_dict = torch.load(custom_weights_path, map_location="cpu", weights_only=True)
+                model.model.load_state_dict(state_dict)
+                print(f"✅ Loaded custom weights from: {custom_weights_path}")
+            else:
+                print("⚠️ No custom weights found, using initialized model")
+                
+        except Exception as e:
+            print(f"⚠️ Could not load custom weights: {e}")
+            print("Using initialized model")
+        
+        # Move model to appropriate device (GPU if available, otherwise CPU)
+        device = kwargs.get("device", None)
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        model = model.to(device)
+        
+        return model
+
+    def save_pretrained(self, save_directory, **kwargs):
+        os.makedirs(save_directory, exist_ok=True)
+        # Save custom weights
+        torch.save(self.model.state_dict(), f"{save_directory}/model.pt")
+        # Save config
+        self.config.save_pretrained(save_directory)
+        # Call parent method for any additional saving
+        super().save_pretrained(save_directory, **kwargs) 

requirements.txt

@@ -0,0 +1,6 @@
+gradio>=4.0.0
+transformers>=4.30.0
+numpy>=1.21.0
+pandas>=1.3.0
+rotary-embedding-torch>=0.2.0
+git+https://github.com/oxpig/ANARCI.git 

restoration.py

@@ -0,0 +1,96 @@
+import numpy as np
+import torch
+
+from extra_utils import res_to_seq, get_sequences_from_anarci
+
+
+class AbRestore:
+    def __init__(self, spread = 11, device = 'cpu', ncpu = 1):
+        self.spread = spread
+        self.device = device
+        self.ncpu = ncpu
+        
+    def _initiate_abrestore(self, model, tokenizer):
+        self.AbLang = model
+        self.tokenizer = tokenizer
+
+    def restore(self, seqs, align = False, **kwargs):
+        """
+        Restore sequences
+        """
+        n_seqs = len(seqs)
+        
+        if align:
+            
+            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:
+            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
+        restored_tokens = torch.where(tokens==23, predicted_tokens, tokens)
+
+        restored_seqs = self.tokenizer(restored_tokens, mode="decode")
+
+        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:])]
+        
+        return np.array([res_to_seq(seq, 'restore') for seq in np.c_[restored_seqs, np.vectorize(len)(seqs)]])
+    
+    def _create_spread_of_sequences(self, seqs, chain = 'H'):
+        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
+        
+        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 _sequence_aligning(self, seqs):
+
+        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)])

scores.py

@@ -0,0 +1,98 @@
+import numpy as np
+import torch
+
+from extra_utils import res_to_list, res_to_seq
+
+
+class AbScores:
+
+    def __init__(self, device = 'cpu', ncpu = 1):
+        
+        self.device = device
+        self.ncpu = ncpu
+        
+    def _initiate_abencoding(self, model, tokenizer):
+        self.AbLang = model
+        self.tokenizer = tokenizer
+        
+    def _encode_sequences(self, seqs):
+        tokens = self.tokenizer(seqs, pad=True, w_extra_tkns=False, device=self.used_device)
+        with torch.no_grad():
+            return self.AbLang.AbRep(tokens).last_hidden_states.numpy()
+        
+    def _predict_logits(self, seqs):
+        tokens = self.tokenizer(seqs, pad=True, w_extra_tkns=False, device=self.used_device)
+        with torch.no_grad():
+            return self.AbLang(tokens), tokens
+        
+    def pseudo_log_likelihood(self, seqs, **kwargs):
+        """
+        Pseudo log likelihood of sequences.
+        """
+        
+        plls = []
+        for seq in seqs:
+            
+            labels = self.tokenizer(
+                seq, pad=True, w_extra_tkns=False, device=self.used_device
+            )
+            
+            idxs = (
+                ~torch.isin(labels, torch.Tensor(self.tokenizer.all_special_tokens).to(self.used_device))
+            ).nonzero()
+            
+            masked_tokens = labels.repeat(len(idxs), 1)
+            for num, idx in enumerate(idxs):
+                masked_tokens[num, idx[1]] = self.tokenizer.mask_token
+            
+            with torch.no_grad():
+                logits = self.AbLang(masked_tokens)
+            
+            logits[:, :, self.tokenizer.all_special_tokens] = -float("inf")            
+            logits = torch.stack([logits[num, idx[1]] for num, idx in enumerate(idxs)])
+  
+            labels = labels[:,idxs[:,1:]].squeeze(2)[0]
+
+            nll = torch.nn.functional.cross_entropy(
+                    logits,
+                    labels,
+                    reduction="mean",
+                )
+            
+            pll = -nll
+
+            plls.append(pll)
+    
+        plls = torch.stack(plls, dim=0).cpu().numpy()
+
+        return plls  
+    
+    def confidence(self, seqs, **kwargs):
+        """
+        Log likelihood of sequences without masking.
+        """
+        
+        labels = self.tokenizer(
+                seqs, pad=True, w_extra_tkns=False, device=self.used_device
+            )
+        with torch.no_grad():
+            logits = self.AbLang(labels)
+            logits[:, :, self.tokenizer.all_special_tokens] = -float("inf")  
+            
+        plls = []
+        for label, logit in zip(labels, logits):
+            
+            idxs = (
+                ~torch.isin(label, torch.Tensor(self.tokenizer.all_special_tokens).to(self.used_device))
+            ).nonzero().squeeze(1)
+
+            nll = torch.nn.functional.cross_entropy(
+                        logit[idxs],
+                        label[idxs],
+                        reduction="mean",
+                    )
+
+            pll = -nll
+            plls.append(pll)
+
+        return torch.stack(plls, dim=0).cpu().numpy()

test_ablang2_HF_implementation.ipynb

@@ -0,0 +1,754 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "458aed0f",
+   "metadata": {},
+   "source": [
+    "<small>Note: This notebook is adapted from the [AbLang2](https://github.com/TobiasHeOl/AbLang2) model's GitHub repository. It is used to verify that the Hugging Face implementation functions correctly and produces the same output as the original model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "a51e7ed2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!rm -rf ~/.cache/huggingface/hub/models--hemantn--ablang2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "7ae54cd0-6253-46dd-a316-4f20b12041e0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "import os\n",
+    "import numpy as np\n",
+    "from transformers import AutoModel, AutoTokenizer\n",
+    "from huggingface_hub import hf_hub_download"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "10801511-770d-46ac-a15d-a02d4ef9ec87",
+   "metadata": {},
+   "source": [
+    "# **0. Sequence input and its format**\n",
+    "\n",
+    "AbLang2 takes as input either the individual heavy variable domain (VH), light variable domain (VL), or the full variable domain (Fv).\n",
+    "\n",
+    "Each record (antibody) needs to be a list with the VH as the first element and the VL as the second. If either the VH or VL is not known, leave an empty string.\n",
+    "\n",
+    "An asterisk (\\*) is used for masking. It is recommended to mask residues which you are interested in mutating.\n",
+    "\n",
+    "**NB:** It is important that the VH and VL sequence is ordered correctly."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "99192978-a008-4a32-a80e-bba238e0ec7c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "seq1 = [\n",
+    "    'EVQLLESGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDVPGHGAAFMDVWGTGTTVTVSS', # VH sequence\n",
+    "    'DIQLTQSPLSLPVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKISNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK' # VL sequence\n",
+    "]\n",
+    "seq2 = [\n",
+    "    'EVQLLESGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDVPGHGAAFMDVWGTGTT',\n",
+    "    'PVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKISNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK'\n",
+    "]\n",
+    "seq3 = [\n",
+    "    'EVQLLESGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDVPGHGAAFMDVWGTGTTVTVSS',\n",
+    "    '' # The VL sequence is not known, so an empty string is left instead. \n",
+    "]\n",
+    "seq4 = [\n",
+    "    '',\n",
+    "    'DIQLTQSPLSLPVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKISNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK'\n",
+    "]\n",
+    "seq5 = [\n",
+    "    'EVQ***SGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCAR**PGHGAAFMDVWGTGTTVTVSS', # (*) is used to mask certain residues\n",
+    "    'DIQLTQSPLSLPVTLGQPASISCRSS*SLEASDTNIYLSWFQQRPGQSPRRLIYKI*NRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK'\n",
+    "]\n",
+    "\n",
+    "all_seqs = [seq1, seq2, seq3, seq4, seq5]\n",
+    "only_both_chains_seqs = [seq1, seq2, seq5]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "dffbacfa-8642-4d94-9572-2205a05c18f9",
+   "metadata": {},
+   "source": [
+    "# **1. How to use AbLang2**\n",
+    "\n",
+    "AbLang2 can be downloaded and used in its raw form as seen below. For convenience, we have also developed different \"modes\" which can be used for specific use cases (see Section 2) "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "6d66ad84",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ed2d5574bd21463c9244070ab762c31e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "config.json:   0%|          | 0.00/763 [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "10e1a02037f74d2da6e0860ef914829b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "configuration_ablang2paired.py: 0.00B [00:00, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "A new version of the following files was downloaded from https://huggingface.co/hemantn/ablang2:\n",
+      "- configuration_ablang2paired.py\n",
+      ". Make sure to double-check they do not contain any added malicious code. To avoid downloading new versions of the code file, you can pin a revision.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "eaf036440107433f950cf4b8c652d756",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "modeling_ablang2paired.py: 0.00B [00:00, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "A new version of the following files was downloaded from https://huggingface.co/hemantn/ablang2:\n",
+      "- modeling_ablang2paired.py\n",
+      ". Make sure to double-check they do not contain any added malicious code. To avoid downloading new versions of the code file, you can pin a revision.\n",
+      "/home/hn533621/.conda/envs/lib_transformer/lib/python3.10/site-packages/huggingface_hub/file_download.py:943: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "22b9a58a3100420c9e353415e7194af6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model.pt:   0%|          | 0.00/179M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "✅ Loaded custom weights from: /home/hn533621/.cache/huggingface/hub/models--hemantn--ablang2/snapshots/13d4401549c368256c517dc13b8ed3d8b28d5e87/model.pt\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e1c40183f9104aa1a67bf9b1c3daea0c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "tokenizer_ablang2paired.py: 0.00B [00:00, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "A new version of the following files was downloaded from https://huggingface.co/hemantn/ablang2:\n",
+      "- tokenizer_ablang2paired.py\n",
+      ". Make sure to double-check they do not contain any added malicious code. To avoid downloading new versions of the code file, you can pin a revision.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3fadab1179e2438ba88e08efb7819680",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "vocab.json:   0%|          | 0.00/331 [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5673cfaa95ac4da78e627c36ad6191b0",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "adapter.py: 0.00B [00:00, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "📁 Files in current directory (/home/hn533621/.cache/huggingface/hub/models--hemantn--ablang2/snapshots/13d4401549c368256c517dc13b8ed3d8b28d5e87):\n",
+      "   adapter.py\n",
+      "   configuration_ablang2paired.py\n",
+      "   tokenizer_ablang2paired.py\n",
+      "   modeling_ablang2paired.py\n",
+      "✅ Successfully imported utility modules from cache directory\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Load model and tokenizer from Hugging Face Hub\n",
+    "model = AutoModel.from_pretrained(\"hemantn/ablang2\", trust_remote_code=True)\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\"hemantn/ablang2\", trust_remote_code=True)\n",
+    "\n",
+    "# Find the cached model directory and import adapter\n",
+    "adapter_path = hf_hub_download(repo_id=\"hemantn/ablang2\", filename=\"adapter.py\")\n",
+    "cached_model_dir = os.path.dirname(adapter_path)\n",
+    "sys.path.insert(0, cached_model_dir)\n",
+    "\n",
+    "# Import and create the adapter\n",
+    "from adapter import AbLang2PairedHuggingFaceAdapter\n",
+    "ablang = AbLang2PairedHuggingFaceAdapter(model=model, tokenizer=tokenizer)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "48562761-6ebe-4025-be97-918c9f9eff7e",
+   "metadata": {},
+   "source": [
+    "# **2. Different modes for specific usecases**\n",
+    "\n",
+    "AbLang2 has already been implemented for a variety of different usecases. The benefit of these modes is that they handle extra tokens such as start, stop and separation tokens.\n",
+    "\n",
+    "1. seqcoding: Generates sequence representations for each sequence\n",
+    "2. rescoding: Generates residue representations for each residue in each sequence\n",
+    "3. likelihood: Generates likelihoods for each amino acid at each position in each sequence\n",
+    "4. probability: Generates probabilities for each amino acid at each position in each sequence\n",
+    "5. pseudo_log_likelihood: Returns the pseudo log likelihood for a sequence (based on masking each residue one at a time)\n",
+    "6. confidence: Returns a fast calculation of the log likelihood for a sequence (based on a single pass with no masking)\n",
+    "7. restore: Restores masked residues\n",
+    "\n",
+    "### **AbLang2 can also align the resulting representations using ANARCI**\n",
+    "\n",
+    "This can be done for 'rescoding', 'likelihood', and 'probability'. This is done by setting the argument \"align=True\".\n",
+    "\n",
+    "**NB**: Align can only be used on input with the same format, i.e. either all heavy, all light, or all both heavy and light.\n",
+    "\n",
+    "### **The align argument can also be used to restore variable missing lengths**\n",
+    "\n",
+    "For this, use \"align=True\" with the 'restore' mode."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "ceae4a88-0679-4704-8bad-c06a4569c497",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "valid_modes = [\n",
+    "    'seqcoding', 'rescoding', 'likelihood', 'probability',\n",
+    "    'pseudo_log_likelihood', 'confidence', 'restore' \n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "aa333732-7508-4826-92ec-3acdd54bc1bb",
+   "metadata": {},
+   "source": [
+    "## **seqcoding** \n",
+    "\n",
+    "The seqcodings represents each sequence as a 480 sized embedding. It is derived from averaging across each rescoding embedding for a given sequence, including extra tokens. \n",
+    "\n",
+    "**NB:** Seqcodings can also be derived in other ways like using the sum or averaging across only parts of the input such as the CDRs. For such cases please use and adapt the below rescoding."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "d22f4302-1262-4cc1-8a1c-a36daa8c710c",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([[-0.2520631 ,  0.18189636,  0.00887137, ...,  0.15365516,\n",
+       "        -0.14508602, -0.13381316],\n",
+       "       [-0.24383117,  0.20946886,  0.07412891, ...,  0.15079288,\n",
+       "        -0.13847049, -0.07304662],\n",
+       "       [-0.20084268,  0.23405147, -0.00103735, ...,  0.07450922,\n",
+       "        -0.08084311, -0.21812904],\n",
+       "       [-0.12659703,  0.3051279 , -0.15117611, ..., -0.20749238,\n",
+       "        -0.10453435, -0.0787883 ],\n",
+       "       [-0.2955319 ,  0.17239201,  0.05676926, ...,  0.15943624,\n",
+       "        -0.16615382, -0.15569784]], shape=(5, 480), dtype=float32)"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ablang(all_seqs, mode='seqcoding')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4b5d9d60",
+   "metadata": {},
+   "source": [
+    "## **rescoding / likelihood / probability**\n",
+    "\n",
+    "The rescodings represents each residue as a 480 sized embedding. The likelihoods represents each residue as the predicted logits for each character in the vocabulary. The probabilities represents the normalised likelihoods.\n",
+    "\n",
+    "**NB:** The output includes extra tokens (start, stop and separation tokens) in the format \"<VH_seq>|<VL_seq>\". The length of the output is therefore 5 longer than the VH and VL.\n",
+    "\n",
+    "**NB:** By default the representations are derived using a single forward pass. To prevent the predicted likelihood and probability to be affected by the input residue at each position, setting the \"stepwise_masking\" argument to True can be used. This will run a forward pass for each position with the residue at that position masked. This is much slower than running a single forward pass."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "6227f661-575f-4b1e-9646-cfba7b10c3b4",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[array([[-0.40741208, -0.5118987 ,  0.06096708, ...,  0.3268144 ,\n",
+       "          0.03920235, -0.36715826],\n",
+       "        [-0.5768883 ,  0.38245413, -0.21791998, ...,  0.01250262,\n",
+       "         -0.08844463, -0.32367525],\n",
+       "        [-0.1475935 ,  0.39639047, -0.38226923, ..., -0.10119921,\n",
+       "         -0.41469565, -0.00319315],\n",
+       "        ...,\n",
+       "        [-0.14358369,  0.3124389 , -0.30157998, ..., -0.13289244,\n",
+       "         -0.45353398, -0.07878865],\n",
+       "        [ 0.17538925,  0.24394299,  0.20141171, ...,  0.14587352,\n",
+       "         -0.38479003,  0.07409196],\n",
+       "        [-0.23031706, -0.35487285,  0.1960684 , ..., -0.1283362 ,\n",
+       "          0.31107333, -0.3265108 ]], shape=(238, 480), dtype=float32),\n",
+       " array([[-0.41981837, -0.3666375 ,  0.10595217, ...,  0.3903574 ,\n",
+       "          0.0382378 , -0.36337993],\n",
+       "        [-0.5054137 ,  0.38347068, -0.10992069, ..., -0.05231472,\n",
+       "         -0.13636623, -0.34830108],\n",
+       "        [-0.06784609,  0.69349885, -0.4212398 , ..., -0.24805346,\n",
+       "         -0.39583805, -0.10972726],\n",
+       "        ...,\n",
+       "        [-0.02212614,  0.26338235, -0.5558968 , ..., -0.24067189,\n",
+       "         -0.11965694,  0.07879876],\n",
+       "        [-0.20650092,  0.43451664, -0.09650223, ..., -0.05296766,\n",
+       "         -0.04297376,  0.41854134],\n",
+       "        [-0.02653179,  0.03729444,  0.13194172, ..., -0.4554279 ,\n",
+       "          0.03723941,  0.17769177]], shape=(238, 480), dtype=float32),\n",
+       " array([[-0.40043733, -0.48596814,  0.0886725 , ...,  0.38941646,\n",
+       "          0.06195956, -0.40999672],\n",
+       "        [-0.54576075,  0.4312959 , -0.3451486 , ..., -0.09285564,\n",
+       "          0.03116508, -0.45269737],\n",
+       "        [ 0.0221165 ,  0.53196615, -0.30137214, ..., -0.1889072 ,\n",
+       "         -0.32587305,  0.05078396],\n",
+       "        ...,\n",
+       "        [-0.03700298,  0.7739084 ,  0.3454928 , ..., -0.03060072,\n",
+       "          0.02420983, -0.48005292],\n",
+       "        [-0.03366657,  0.74771184, -0.35423476, ..., -0.08759108,\n",
+       "         -0.17898935, -0.4540483 ],\n",
+       "        [-0.16625853,  0.2701079 , -0.19761363, ...,  0.10313392,\n",
+       "          0.44890267, -0.64840287]], shape=(238, 480), dtype=float32),\n",
+       " array([[-0.26863217,  0.32259187,  0.10813517, ...,  0.03953876,\n",
+       "          0.18312076, -0.00498045],\n",
+       "        [-0.2165424 , -0.38562432, -0.02696264, ...,  0.20541488,\n",
+       "          0.18698391, -0.22639504],\n",
+       "        [-0.41950518,  0.04743317,  0.0048816 , ...,  0.11408642,\n",
+       "         -0.05384652,  0.1025871 ],\n",
+       "        ...,\n",
+       "        [-0.14095458,  0.5860325 , -0.44657114, ..., -0.39150292,\n",
+       "         -0.22395667, -0.42516366],\n",
+       "        [ 0.29816052,  0.40440455, -0.52062094, ...,  0.08969188,\n",
+       "         -0.20792632, -0.2045222 ],\n",
+       "        [-0.21370608,  0.23035707, -0.355185  , ..., -0.36726946,\n",
+       "         -0.05693531, -0.37847823]], shape=(238, 480), dtype=float32),\n",
+       " array([[-0.42062947, -0.44009134,  0.00152371, ...,  0.27141467,\n",
+       "          0.03798106, -0.397461  ],\n",
+       "        [-0.57318133,  0.5258899 , -0.17001636, ..., -0.23864633,\n",
+       "          0.2088059 , -0.57877594],\n",
+       "        [-0.38988614,  0.46168196, -0.3429413 , ..., -0.14872643,\n",
+       "         -0.46576905, -0.21224979],\n",
+       "        ...,\n",
+       "        [-0.21528634,  0.30046722, -0.25216463, ..., -0.11576828,\n",
+       "         -0.4704907 , -0.0740136 ],\n",
+       "        [ 0.0633081 ,  0.22700705,  0.28184187, ...,  0.15967266,\n",
+       "         -0.377182  ,  0.06188517],\n",
+       "        [-0.27826303, -0.37297496,  0.21229912, ..., -0.14886017,\n",
+       "          0.24998347, -0.35954213]], shape=(238, 480), dtype=float32)]"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ablang(all_seqs, mode='rescoding', stepwise_masking = False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6da2183b-4306-49bd-a7fc-23e78a23f305",
+   "metadata": {},
+   "source": [
+    "## **Align rescoding/likelihood/probability output**\n",
+    "\n",
+    "For the 'rescoding', 'likelihood', and 'probability' modes, the output can also be aligned using the argument \"align=True\".\n",
+    "\n",
+    "This is done using the antibody numbering tool ANARCI, and requires manually installing **Pandas** and **[ANARCI](https://github.com/oxpig/ANARCI)**.\n",
+    "\n",
+    "**NB**: Align can only be used on input with the same format, i.e. either all heavy, all light, or all both heavy and light."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "e4bc0cb1-f5b0-4255-9e93-d643ae1396df",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "['<' '1 ' '2 ' '3 ' '4 ' '5 ' '6 ' '7 ' '8 ' '9 ' '11 ' '12 ' '13 ' '14 '\n",
+      " '15 ' '16 ' '17 ' '18 ' '19 ' '20 ' '21 ' '22 ' '23 ' '24 ' '25 ' '26 '\n",
+      " '27 ' '28 ' '29 ' '30 ' '35 ' '36 ' '37 ' '38 ' '39 ' '40 ' '41 ' '42 '\n",
+      " '43 ' '44 ' '45 ' '46 ' '47 ' '48 ' '49 ' '50 ' '51 ' '52 ' '53 ' '54 '\n",
+      " '55 ' '56 ' '57 ' '58 ' '59 ' '62 ' '63 ' '64 ' '65 ' '66 ' '67 ' '68 '\n",
+      " '69 ' '70 ' '71 ' '72 ' '74 ' '75 ' '76 ' '77 ' '78 ' '79 ' '80 ' '81 '\n",
+      " '82 ' '83 ' '84 ' '85 ' '86 ' '87 ' '88 ' '89 ' '90 ' '91 ' '92 ' '93 '\n",
+      " '94 ' '95 ' '96 ' '97 ' '98 ' '99 ' '100 ' '101 ' '102 ' '103 ' '104 '\n",
+      " '105 ' '106 ' '107 ' '108 ' '109 ' '110 ' '111 ' '112A' '112 ' '113 '\n",
+      " '114 ' '115 ' '116 ' '117 ' '118 ' '119 ' '120 ' '121 ' '122 ' '123 '\n",
+      " '124 ' '125 ' '126 ' '127 ' '128 ' '>' '|' '<' '1 ' '2 ' '3 ' '4 ' '5 '\n",
+      " '6 ' '7 ' '8 ' '9 ' '10 ' '11 ' '12 ' '13 ' '14 ' '15 ' '16 ' '17 ' '18 '\n",
+      " '19 ' '20 ' '21 ' '22 ' '23 ' '24 ' '25 ' '26 ' '27 ' '28 ' '29 ' '30 '\n",
+      " '31 ' '32 ' '34 ' '35 ' '36 ' '37 ' '38 ' '39 ' '40 ' '41 ' '42 ' '43 '\n",
+      " '44 ' '45 ' '46 ' '47 ' '48 ' '49 ' '50 ' '51 ' '52 ' '53 ' '54 ' '55 '\n",
+      " '56 ' '57 ' '64 ' '65 ' '66 ' '67 ' '68 ' '69 ' '70 ' '71 ' '72 ' '74 '\n",
+      " '75 ' '76 ' '77 ' '78 ' '79 ' '80 ' '83 ' '84 ' '85 ' '86 ' '87 ' '88 '\n",
+      " '89 ' '90 ' '91 ' '92 ' '93 ' '94 ' '95 ' '96 ' '97 ' '98 ' '99 ' '100 '\n",
+      " '101 ' '102 ' '103 ' '104 ' '105 ' '106 ' '107 ' '108 ' '109 ' '114 '\n",
+      " '115 ' '116 ' '117 ' '118 ' '119 ' '120 ' '121 ' '122 ' '123 ' '124 '\n",
+      " '125 ' '126 ' '127 ' '>']\n",
+      "['<EVQLLESGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDVPGHGAAFMDVWGTGTTVTVSS>|<DIQLTQSPLSLPVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKI-SNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK>', '<EVQLLESGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDVPGHGAAFMDVWGTGTT----->|<-----------PVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKI-SNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK>', '<------SGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCAR**PGHGAAFMDVWGTGTTVTVSS>|<DIQLTQSPLSLPVTLGQPASISCRSS*SLEASDTNIYLSWFQQRPGQSPRRLIYKI*N-RDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK>']\n",
+      "[[[  9.31621838  -3.42184329  -3.59397745 ... -14.73707485  -6.8935833\n",
+      "    -0.23662776]\n",
+      "  [ -3.54718232  -5.84866619  -4.02423859 ... -12.93966579  -9.5614481\n",
+      "    -4.48473835]\n",
+      "  [-11.94997597  -2.245543    -5.69481373 ... -15.19639015 -17.97454071\n",
+      "   -12.56952095]\n",
+      "  ...\n",
+      "  [ -8.94504833  -0.42261261  -4.95588207 ... -16.66817474 -15.2224741\n",
+      "   -10.37267494]\n",
+      "  [-11.65150356  -5.44477606  -2.95585775 ... -16.25555801  -9.75158596\n",
+      "   -11.75897026]\n",
+      "  [  1.79469728  -1.95846701  -3.59784532 ... -14.95585823  -7.47080708\n",
+      "    -0.95226753]]\n",
+      "\n",
+      " [[  8.55518723  -3.83663297  -2.33595967 ... -13.87456799  -8.14840603\n",
+      "    -0.42472434]\n",
+      "  [ -4.40701294  -5.53201008  -3.69397402 ... -12.97877789  -9.86258411\n",
+      "    -4.95414352]\n",
+      "  [-11.95642853  -3.86210871  -5.80935192 ... -14.89213085 -16.94556236\n",
+      "   -11.36959839]\n",
+      "  ...\n",
+      "  [ -7.75924015  -0.66524202  -4.08643246 ... -16.16580772 -14.76507473\n",
+      "    -8.3507061 ]\n",
+      "  [-11.91039753  -4.86995983  -2.74777436 ... -16.07694817  -8.44974899\n",
+      "   -10.45223904]\n",
+      "  [  0.86006832  -2.37964034  -3.58130741 ... -15.35423565  -7.73035526\n",
+      "    -1.11989737]]\n",
+      "\n",
+      " [[ -4.37902737  -7.55587149   1.21958363 ... -15.48622513  -6.021842\n",
+      "    -3.79647374]\n",
+      "  [  0.           0.           0.         ...   0.           0.\n",
+      "     0.        ]\n",
+      "  [  0.           0.           0.         ...   0.           0.\n",
+      "     0.        ]\n",
+      "  ...\n",
+      "  [ -8.94207573  -0.51090252  -5.09760332 ... -16.69521713 -15.45450687\n",
+      "   -10.50823212]\n",
+      "  [-11.92354965  -5.55152607  -2.87666893 ... -16.40607834 -10.19431686\n",
+      "   -12.1328764 ]\n",
+      "  [  2.42200375  -2.01573253  -3.61701298 ... -14.9590435   -7.19029331\n",
+      "    -0.89830256]]]\n"
+     ]
+    }
+   ],
+   "source": [
+    "results = ablang(only_both_chains_seqs, mode='likelihood', align=True)\n",
+    "\n",
+    "print(results.number_alignment)\n",
+    "print(results.aligned_seqs)\n",
+    "print(results.aligned_embeds)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "56be8cad",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[array([[9.9955505e-01, 2.9358694e-06, 2.4716087e-06, ..., 3.5776201e-11,\n",
+       "         9.1196831e-08, 7.0967326e-05],\n",
+       "        [4.1573694e-06, 4.1619489e-07, 2.5800944e-06, ..., 3.4650952e-10,\n",
+       "         1.0159109e-08, 1.6279575e-06],\n",
+       "        [7.8059600e-08, 1.2794037e-03, 4.0645118e-05, ..., 3.0375720e-09,\n",
+       "         1.8879491e-10, 4.2010839e-08],\n",
+       "        ...,\n",
+       "        [3.4210879e-07, 1.7195340e-03, 1.8477240e-05, ..., 1.5137445e-10,\n",
+       "         6.4255873e-10, 8.2064140e-08],\n",
+       "        [9.1038084e-09, 4.5161755e-06, 5.4411950e-05, ..., 9.1139631e-11,\n",
+       "         6.0862085e-08, 8.1761966e-09],\n",
+       "        [8.5759175e-04, 2.0104915e-05, 3.9023766e-06, ..., 4.5562460e-11,\n",
+       "         8.1156479e-08, 5.4990651e-05]], shape=(238, 26), dtype=float32),\n",
+       " array([[9.9939799e-01, 4.1499175e-06, 1.8611167e-05, ..., 1.8139243e-10,\n",
+       "         5.5649299e-08, 1.2583815e-04],\n",
+       "        [1.6735513e-06, 5.4332406e-07, 3.4143472e-06, ..., 3.1693398e-10,\n",
+       "         7.1501400e-09, 9.6832969e-07],\n",
+       "        [3.7784993e-08, 1.2377645e-04, 1.7658784e-05, ..., 2.0061326e-09,\n",
+       "         2.5737484e-10, 6.7947965e-08],\n",
+       "        ...,\n",
+       "        [1.1050455e-06, 1.3312638e-03, 4.3497097e-05, ..., 2.4686178e-10,\n",
+       "         1.0018089e-09, 6.1165900e-07],\n",
+       "        [5.7270397e-09, 6.5396339e-06, 5.4601755e-05, ..., 8.8801404e-11,\n",
+       "         1.8233513e-07, 2.4615032e-08],\n",
+       "        [7.3952030e-04, 2.8970928e-05, 8.7113440e-06, ..., 6.7168833e-11,\n",
+       "         1.3746008e-07, 1.0210846e-04]], shape=(222, 26), dtype=float32),\n",
+       " array([[9.99685407e-01, 3.35662639e-06, 1.14241482e-06, ...,\n",
+       "         2.32460891e-11, 6.88188067e-08, 5.69467156e-05],\n",
+       "        [6.38133372e-07, 1.01300586e-07, 5.64459742e-06, ...,\n",
+       "         4.09234556e-11, 2.53804799e-09, 4.31722100e-07],\n",
+       "        [1.49096788e-08, 2.04515047e-04, 9.23794141e-06, ...,\n",
+       "         7.46306961e-10, 2.92107380e-11, 2.21786500e-08],\n",
+       "        ...,\n",
+       "        [2.15093763e-07, 1.06453872e-03, 1.62486140e-05, ...,\n",
+       "         1.12102910e-10, 1.47300866e-10, 4.73037538e-08],\n",
+       "        [4.30136682e-09, 3.09317988e-06, 3.96632568e-05, ...,\n",
+       "         5.24226877e-11, 2.39579450e-08, 3.86403221e-09],\n",
+       "        [9.77773685e-04, 1.29533228e-05, 2.78623725e-06, ...,\n",
+       "         2.73364300e-11, 3.96418649e-08, 4.04014427e-05]],\n",
+       "       shape=(238, 26), dtype=float32)]"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ablang(only_both_chains_seqs, mode='probability')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8f0a71ec-e916-4330-90d0-13a4b1121a89",
+   "metadata": {},
+   "source": [
+    "## **Pseudo log likelihood and Confidence scores**\n",
+    "\n",
+    "The pseudo log likelihood and confidence represents two methods for calculating the uncertainty for the input sequence.\n",
+    "\n",
+    "- pseudo_log_likelihood: For each position, the pseudo log likelihood is calculated when predicting the masked residue. The final score is an average across the whole input. This is similar to the approach taken in the ESM-2 paper for calculating pseudo perplexity [(Lin et al., 2023)](https://doi.org/10.1126/science.ade2574).\n",
+    "\n",
+    "- confidence: For each position, the log likelihood is calculated without masking the residue. The final score is an average across the whole input. \n",
+    "\n",
+    "**NB:** The **confidence is fast** to compute, requiring only a single forward pass per input. **Pseudo log likelihood is slow** to calculate, requiring L forward passes per input, where L is the length of the input.\n",
+    "\n",
+    "**NB:** It is recommended to use **pseudo log likelihood for final results** and **confidence for exploratory work**."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "83f3064b-48a7-42fb-ba82-ec153ea946da",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([1.96673731, 2.04801253, 2.09881898, 1.82533665, 1.97255249])"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "results = ablang(all_seqs, mode='pseudo_log_likelihood')\n",
+    "np.exp(-results) # convert to pseudo perplexity"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "42cc8b34-5ae9-4857-93fe-a438a0f2a868",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([1.2636038, 1.126463 , 1.3123759, 1.2140924, 1.1805094],\n",
+       "      dtype=float32)"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "results = ablang(all_seqs, mode='confidence')\n",
+    "np.exp(-results)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e0b63e48-b2a1-4a8e-8ecb-449748a2cb25",
+   "metadata": {},
+   "source": [
+    "## **restore**\n",
+    "\n",
+    "This mode can be used to restore masked residues, and fragmented regions with \"align=True\". "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "2d5b725c-4eac-4a4b-9331-357c3ac140f7",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array(['<EVQLLESGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDVPGHGAAFMDVWGTGTTVTVSS>|<DIQLTQSPLSLPVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKISNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK>',\n",
+       "       '<EVQLLESGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDVPGHGAAFMDVWGTGTT>|<PVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKISNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK>',\n",
+       "       '<EVQLVQSGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDPPGHGAAFMDVWGTGTTVTVSS>|<DIQLTQSPLSLPVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKISNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK>'],\n",
+       "      dtype='<U238')"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "restored = ablang(only_both_chains_seqs, mode='restore')\n",
+    "restored"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "0e9615f7-c490-4947-96f4-7617266c686e",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array(['<EVQLLESGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDVPGHGAAFMDVWGTGTTVTVSS>|<DIQLTQSPLSLPVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKISNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK>',\n",
+       "       '<EVQLLESGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDVPGHGAAFMDVWGTGTTVTVSS>|<DVVMTQSPLSLPVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKISNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK>',\n",
+       "       '<QVQLVQSGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCARDPPGHGAAFMDVWGTGTTVTVSS>|<DIQLTQSPLSLPVTLGQPASISCRSSQSLEASDTNIYLSWFQQRPGQSPRRLIYKISNRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK>'],\n",
+       "      dtype='<U238')"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "restored = ablang(only_both_chains_seqs, mode='restore', align = True)\n",
+    "restored"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d80020ce",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "lib_transformer",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.18"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

tokenizer_ablang2paired.py

@@ -0,0 +1,124 @@
+import json
+import os
+from transformers import PreTrainedTokenizer
+
+
+class AbLang2PairedTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "vocab.json"}
+    model_input_names = ["input_ids"]
+
+    def __init__(self, vocab_file=None, **kwargs):
+        if vocab_file is None:
+            # Try to find vocab file in the current directory
+            vocab_file = "vocab.json"
+        
+        self.vocab_file = vocab_file
+        with open(vocab_file, "r", encoding="utf-8") as f:
+            self.vocab = json.load(f)
+        
+        # Set required token attributes (all as strings, standard for HF)
+        kwargs.setdefault("pad_token", "-")
+        kwargs.setdefault("mask_token", "*")
+        kwargs.setdefault("unk_token", "X")
+        
+        super().__init__(**kwargs)
+
+    @property
+    def pad_token_id(self):
+        return self.vocab[self.pad_token]
+
+    @property
+    def mask_token_id(self):
+        return self.vocab[self.mask_token]
+
+    def _tokenize(self, text):
+        return list(text)
+    
+    def tokenize(self, text, text_pair=None, **kwargs):
+        """Tokenize text or text pair."""
+        if text_pair is not None:
+            # For paired sequences, combine them with a separator
+            combined_text = text + "|" + text_pair
+            return self._tokenize(combined_text)
+        else:
+            return self._tokenize(text)
+
+    def _convert_token_to_id(self, token):
+        return self.vocab.get(token, self.vocab[self.unk_token])
+
+    def _convert_id_to_token(self, index):
+        inv_vocab = {v: k for k, v in self.vocab.items()}
+        return inv_vocab.get(index, self.unk_token)
+
+    def get_vocab(self):
+        return self.vocab
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        os.makedirs(save_directory, exist_ok=True)
+        path = os.path.join(save_directory, (filename_prefix or "") + "vocab.json")
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(self.vocab, f)
+        return (path,)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        try:
+            from transformers.utils import cached_file
+            vocab_file = cached_file(
+                pretrained_model_name_or_path,
+                "vocab.json",
+                cache_dir=kwargs.get("cache_dir"),
+                force_download=kwargs.get("force_download", False),
+                resume_download=kwargs.get("resume_download", False),
+                proxies=kwargs.get("proxies"),
+                token=kwargs.get("token"),
+                revision=kwargs.get("revision"),
+                local_files_only=kwargs.get("local_files_only", False),
+            )
+            
+            if vocab_file is None or not os.path.exists(vocab_file):
+                raise ValueError(f"Vocabulary file vocab.json not found in {pretrained_model_name_or_path}")
+                
+            return cls(vocab_file=vocab_file, **kwargs)
+        except Exception as e:
+            # Fallback for local paths
+            vocab_file = os.path.join(pretrained_model_name_or_path, "vocab.json")
+            if not os.path.exists(vocab_file):
+                raise ValueError(f"Vocabulary file {vocab_file} not found")
+            return cls(vocab_file=vocab_file, **kwargs)
+
+    def save_pretrained(self, save_directory, filename_prefix=None):
+        os.makedirs(save_directory, exist_ok=True)
+        vocab_files = self.save_vocabulary(save_directory, filename_prefix)
+
+        tokenizer_config = {
+            "tokenizer_class": f"{self.__class__.__module__}.{self.__class__.__name__}"
+        }
+        with open(os.path.join(save_directory, "tokenizer_config.json"), "w") as f:
+            json.dump(tokenizer_config, f, indent=2)
+
+        return vocab_files
+
+    def __call__(self, sequences, padding=False, return_tensors=None, **kwargs):
+        # Handle different input formats
+        if isinstance(sequences, str):
+            # Single string: "VH|VL"
+            sequences = [sequences]
+        elif isinstance(sequences, list) and len(sequences) > 0:
+            if isinstance(sequences[0], list):
+                # List of lists: [['VH', 'VL'], ['VH2', 'VL2']]
+                sequences = [f"{pair[0]}|{pair[1]}" for pair in sequences]
+            # List of strings: ["VH|VL", "VH2|VL2"] - already correct format
+        
+        # Tokenize each sequence
+        input_ids = [[self._convert_token_to_id(tok) for tok in self._tokenize(seq)] for seq in sequences]
+        # Padding
+        if padding:
+            maxlen = max(len(ids) for ids in input_ids)
+            input_ids = [ids + [self.pad_token_id] * (maxlen - len(ids)) for ids in input_ids]
+        # Return tensors if requested
+        if return_tensors == 'pt':
+            import torch
+            input_ids = torch.tensor(input_ids)
+        return {'input_ids': input_ids}
+

vocab.json

@@ -0,0 +1,28 @@
+{
+    "<": 0,
+    "M": 1,
+    "R": 2,
+    "H": 3,
+    "K": 4,
+    "D": 5,
+    "E": 6,
+    "S": 7,
+    "T": 8,
+    "N": 9,
+    "Q": 10,
+    "C": 11,
+    "G": 12,
+    "P": 13,
+    "A": 14,
+    "V": 15,
+    "I": 16,
+    "F": 17,
+    "Y": 18,
+    "W": 19,
+    "L": 20,
+    "-": 21,
+    ">": 22,
+    "*": 23,
+    "X": 24,
+    "|": 25
+} 

vocab.py

@@ -0,0 +1,28 @@
+ablang_vocab = {
+    "<": 0,
+    "M": 1,
+    "R": 2,
+    "H": 3,
+    "K": 4,
+    "D": 5,
+    "E": 6,
+    "S": 7,
+    "T": 8,
+    "N": 9,
+    "Q": 10,
+    "C": 11,
+    "G": 12,
+    "P": 13,
+    "A": 14,
+    "V": 15,
+    "I": 16,
+    "F": 17,
+    "Y": 18,
+    "W": 19,
+    "L": 20,
+    "-": 21,
+    ">": 22,
+    "*": 23,
+    "X": 24,
+    "|": 25
+}