Update

README.md

@@ -0,0 +1,14 @@
+---
+title: Test Webpage
+emoji: 🐢
+colorFrom: blue
+colorTo: blue
+sdk: gradio
+sdk_version: 5.7.1
+app_file: app.py
+pinned: false
+license: mit
+short_description: test_webpage
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,653 @@
+from datetime import datetime
+import gradio as gr
+import requests
+from Bio.PDB import PDBParser, MMCIFParser, PDBIO, Select
+from Bio.PDB.Polypeptide import is_aa
+from Bio.SeqUtils import seq1
+from typing import Optional, Tuple
+import numpy as np
+import os
+from gradio_molecule3d import Molecule3D
+
+from model_loader import load_model
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+import re
+import pandas as pd
+import copy
+
+import transformers
+from transformers import AutoTokenizer, DataCollatorForTokenClassification
+
+from datasets import Dataset
+
+from scipy.special import expit
+
+# Load model and move to device
+#checkpoint = 'ThorbenF/prot_t5_xl_uniref50'
+#checkpoint = 'ThorbenF/prot_t5_xl_uniref50_cryptic'
+#checkpoint = 'ThorbenF/prot_t5_xl_uniref50_database'
+#checkpoint = 'ThorbenF/prot_t5_xl_uniref50_full'
+#checkpoint = 'ThorbenF/prot_t5_xl_uniref50_0925'
+#checkpoint = 'ThorbenF/prot_t5_xl_uniref50_0925_v2'
+checkpoint = 'ThorbenF/prot_t5_xl_uniref50_full_v2'
+max_length = 1500
+model, tokenizer = load_model(checkpoint, max_length)
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model.to(device)
+model.eval()
+
+def normalize_scores(scores):
+    min_score = np.min(scores)
+    max_score = np.max(scores)
+    return (scores - min_score) / (max_score - min_score) if max_score > min_score else scores
+
+def read_mol(pdb_path):
+    """Read PDB file and return its content as a string"""
+    with open(pdb_path, 'r') as f:
+        return f.read()
+
+def fetch_structure(pdb_id: str, output_dir: str = ".") -> str:
+    """
+    Fetch the structure file for a given PDB ID. Prioritizes CIF files.
+    If a structure file already exists locally, it uses that.
+    """
+    file_path = download_structure(pdb_id, output_dir)
+    return file_path
+
+def download_structure(pdb_id: str, output_dir: str) -> str:
+    """
+    Attempt to download the structure file in CIF or PDB format.
+    Returns the path to the downloaded file.
+    """
+    for ext in ['.cif', '.pdb']:
+        file_path = os.path.join(output_dir, f"{pdb_id}{ext}")
+        if os.path.exists(file_path):
+            return file_path
+        url = f"https://files.rcsb.org/download/{pdb_id}{ext}"
+        response = requests.get(url, timeout=10)
+        if response.status_code == 200:
+            with open(file_path, 'wb') as f:
+                f.write(response.content)
+            return file_path
+    return None
+
+def convert_cif_to_pdb(cif_path: str, output_dir: str = ".") -> str:
+    """
+    Convert a CIF file to PDB format using BioPython and return the PDB file path.
+    """
+    pdb_path = os.path.join(output_dir, os.path.basename(cif_path).replace('.cif', '.pdb'))
+    parser = MMCIFParser(QUIET=True)
+    structure = parser.get_structure('protein', cif_path)
+    io = PDBIO()
+    io.set_structure(structure)
+    io.save(pdb_path)
+    return pdb_path
+
+def fetch_pdb(pdb_id):
+    pdb_path = fetch_structure(pdb_id)
+    _, ext = os.path.splitext(pdb_path)
+    if ext == '.cif':
+        pdb_path = convert_cif_to_pdb(pdb_path)
+    return pdb_path
+
+def create_chain_specific_pdb(input_pdb: str, chain_id: str, residue_scores: list, protein_residues: list) -> str:
+    """
+    Create a PDB file with only the selected chain and residues, replacing B-factor with prediction scores
+    """
+    parser = PDBParser(QUIET=True)
+    structure = parser.get_structure('protein', input_pdb)
+    
+    output_pdb = f"{os.path.splitext(input_pdb)[0]}_{chain_id}_predictions_scores.pdb"
+    
+    # Create scores dictionary for easy lookup
+    scores_dict = {resi: score for resi, score in residue_scores}
+
+    # Create a custom Select class
+    class ResidueSelector(Select):
+        def __init__(self, chain_id, selected_residues, scores_dict):
+            self.chain_id = chain_id
+            self.selected_residues = selected_residues
+            self.scores_dict = scores_dict
+        
+        def accept_chain(self, chain):
+            return chain.id == self.chain_id
+        
+        def accept_residue(self, residue):
+            return residue.id[1] in self.selected_residues
+
+        def accept_atom(self, atom):
+            if atom.parent.id[1] in self.scores_dict:
+                atom.bfactor = np.absolute(1-self.scores_dict[atom.parent.id[1]]) * 100
+            return True
+
+    # Prepare output PDB with selected chain and residues, modified B-factors
+    io = PDBIO()
+    selector = ResidueSelector(chain_id, [res.id[1] for res in protein_residues], scores_dict)
+    
+    io.set_structure(structure[0])
+    io.save(output_pdb, selector)
+    
+    return output_pdb
+
+def generate_pymol_commands(pdb_id, segment, residues_by_bracket, current_time, score_type):
+    """Generate PyMOL commands based on score type"""
+    pymol_commands = f"Prediction for PDB: {pdb_id}, Chain: {segment}\nDate: {current_time}\nScore Type: {score_type}\n\n"
+    
+    pymol_commands += f"""
+    # PyMOL Visualization Commands
+    fetch {pdb_id}, protein
+    hide everything, all
+    show cartoon, chain {segment}
+    color white, chain {segment}
+    """
+    
+    # Define colors for each score bracket
+    bracket_colors = {
+        "0.0-0.2": "white",
+        "0.2-0.4": "lightorange",
+        "0.4-0.6": "yelloworange",
+        "0.6-0.8": "orange",
+        "0.8-1.0": "red"
+    }
+    
+    # Add PyMOL commands for each score bracket
+    for bracket, residues in residues_by_bracket.items():
+        if residues:  # Only add commands if there are residues in this bracket
+            color = bracket_colors[bracket]
+            resi_list = '+'.join(map(str, residues))
+            pymol_commands += f"""
+    select bracket_{bracket.replace('.', '').replace('-', '_')}, resi {resi_list} and chain {segment}
+    show sticks, bracket_{bracket.replace('.', '').replace('-', '_')}
+    color {color}, bracket_{bracket.replace('.', '').replace('-', '_')}
+    """
+    return pymol_commands
+
+def generate_results_text(pdb_id, segment, residues_by_bracket, protein_residues, sequence, scores, current_time, score_type):
+    """Generate results text based on score type"""
+    result_str = f"Prediction for PDB: {pdb_id}, Chain: {segment}\nDate: {current_time}\nScore Type: {score_type}\n\n"
+    result_str += "Residues by Score Brackets:\n\n"
+    
+    # Add residues for each bracket
+    for bracket, residues in residues_by_bracket.items():
+        result_str += f"Bracket {bracket}:\n"
+        result_str += f"Columns: Residue Name, Residue Number, One-letter Code, {score_type} Score\n"
+        result_str += "\n".join([
+            f"{res.resname} {res.id[1]} {sequence[i]} {scores[i]:.2f}" 
+            for i, res in enumerate(protein_residues) if res.id[1] in residues
+        ])
+        result_str += "\n\n"
+    
+    return result_str
+
+
+
+
+def process_pdb(pdb_id_or_file, segment, score_type='normalized'):
+    # Determine if input is a PDB ID or file path
+    if pdb_id_or_file.endswith('.pdb'):
+        pdb_path = pdb_id_or_file
+        pdb_id = os.path.splitext(os.path.basename(pdb_path))[0]
+    else:
+        pdb_id = pdb_id_or_file
+        pdb_path = fetch_pdb(pdb_id)
+    
+    # Determine the file format and choose the appropriate parser
+    _, ext = os.path.splitext(pdb_path)
+    parser = MMCIFParser(QUIET=True) if ext == '.cif' else PDBParser(QUIET=True)
+    
+    # Parse the structure file
+    structure = parser.get_structure('protein', pdb_path)
+    
+    # Extract the specified chain
+    chain = structure[0][segment]
+    
+    protein_residues = [res for res in chain if is_aa(res)]
+    sequence = "".join(seq1(res.resname) for res in protein_residues)
+    sequence_id = [res.id[1] for res in protein_residues]
+
+    input_ids = tokenizer(" ".join(sequence), return_tensors="pt").input_ids.to(device)
+    with torch.no_grad():
+        outputs = model(input_ids).logits.detach().cpu().numpy().squeeze()
+
+    # Calculate scores and normalize them
+    raw_scores = expit(outputs[:, 1] - outputs[:, 0])
+    normalized_scores = normalize_scores(raw_scores)
+    
+    # Choose which scores to use based on score_type
+    display_scores = normalized_scores if score_type == 'normalized' else raw_scores
+    
+    # Zip residues with scores to track the residue ID and score
+    residue_scores = [(resi, score) for resi, score in zip(sequence_id, display_scores)]
+    
+    # Also save both score types for later use
+    raw_residue_scores = [(resi, score) for resi, score in zip(sequence_id, raw_scores)]
+    norm_residue_scores = [(resi, score) for resi, score in zip(sequence_id, normalized_scores)]
+    
+    # Define the score brackets
+    score_brackets = {
+        "0.0-0.2": (0.0, 0.2),
+        "0.2-0.4": (0.2, 0.4),
+        "0.4-0.6": (0.4, 0.6),
+        "0.6-0.8": (0.6, 0.8),
+        "0.8-1.0": (0.8, 1.0)
+    }
+    
+    # Initialize a dictionary to store residues by bracket
+    residues_by_bracket = {bracket: [] for bracket in score_brackets}
+    
+    # Categorize residues into brackets
+    for resi, score in residue_scores:
+        for bracket, (lower, upper) in score_brackets.items():
+            if lower <= score < upper:
+                residues_by_bracket[bracket].append(resi)
+                break
+    
+    # Generate timestamp
+    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    
+    # Generate result text and PyMOL commands based on score type
+    display_score_type = "Normalized" if score_type == 'normalized' else "Raw"
+    result_str = generate_results_text(pdb_id, segment, residues_by_bracket, protein_residues, sequence, 
+                                      display_scores, current_time, display_score_type)
+    pymol_commands = generate_pymol_commands(pdb_id, segment, residues_by_bracket, current_time, display_score_type)
+    
+    # Create chain-specific PDB with scores in B-factor
+    scored_pdb = create_chain_specific_pdb(pdb_path, segment, residue_scores, protein_residues)
+
+    # Molecule visualization with updated script with color mapping
+    mol_vis = molecule(pdb_path, residue_scores, segment)
+    
+    # Create prediction file
+    prediction_file = f"{pdb_id}_{display_score_type.lower()}_binding_site_residues.txt"
+    with open(prediction_file, "w") as f:
+        f.write(result_str)
+    
+    scored_pdb_name = f"{pdb_id}_{segment}_{display_score_type.lower()}_predictions_scores.pdb"
+    os.rename(scored_pdb, scored_pdb_name)
+    
+    return pymol_commands, mol_vis, [prediction_file, scored_pdb_name], raw_residue_scores, norm_residue_scores, pdb_id, segment
+
+def molecule(input_pdb, residue_scores=None, segment='A'):
+    # Read PDB file content
+    mol = read_mol(input_pdb)
+    
+    # Prepare high-scoring residues script if scores are provided
+    high_score_script = ""
+    if residue_scores is not None:
+        # Filter residues based on their scores
+        class1_score_residues = [resi for resi, score in residue_scores if 0.0 < score <= 0.2]
+        class2_score_residues = [resi for resi, score in residue_scores if 0.2 < score <= 0.4]
+        class3_score_residues = [resi for resi, score in residue_scores if 0.4 < score <= 0.6]
+        class4_score_residues = [resi for resi, score in residue_scores if 0.6 < score <= 0.8]
+        class5_score_residues = [resi for resi, score in residue_scores if 0.8 < score <= 1.0]
+        
+        high_score_script = """
+        // Load the original model and apply white cartoon style
+        let chainModel = viewer.addModel(pdb, "pdb");
+        chainModel.setStyle({}, {});
+        chainModel.setStyle(
+            {"chain": "%s"}, 
+            {"cartoon": {"color": "white"}}
+        );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class1Model = viewer.addModel(pdb, "pdb");
+        class1Model.setStyle({}, {});
+        class1Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "0xFFFFFF", "opacity": 0.5}}
+        );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class2Model = viewer.addModel(pdb, "pdb");
+        class2Model.setStyle({}, {});
+        class2Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "0xFFD580", "opacity": 0.7}}
+        );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class3Model = viewer.addModel(pdb, "pdb");
+        class3Model.setStyle({}, {});
+        class3Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "0xFFA500", "opacity": 1}}
+        );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class4Model = viewer.addModel(pdb, "pdb");
+        class4Model.setStyle({}, {});
+        class4Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "0xFF4500", "opacity": 1}}
+        );
+
+        // Create a new model for high-scoring residues and apply red sticks style
+        let class5Model = viewer.addModel(pdb, "pdb");
+        class5Model.setStyle({}, {});
+        class5Model.setStyle(
+            {"chain": "%s", "resi": [%s]}, 
+            {"stick": {"color": "0xFF0000", "alpha": 1}}
+        );
+
+        """ % (
+            segment,
+            segment,
+            ", ".join(str(resi) for resi in class1_score_residues),
+            segment,
+            ", ".join(str(resi) for resi in class2_score_residues),
+            segment,
+            ", ".join(str(resi) for resi in class3_score_residues),
+            segment,
+            ", ".join(str(resi) for resi in class4_score_residues),
+            segment,
+            ", ".join(str(resi) for resi in class5_score_residues)
+        )
+    
+    # Generate the full HTML content
+    html_content = f"""
+    <!DOCTYPE html>
+    <html>
+    <head>    
+        <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
+        <style>
+        .mol-container {{
+            width: 100%;
+            height: 700px;
+            position: relative;
+        }}
+        </style>
+        <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js"></script>
+        <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
+    </head>
+    <body>
+        <div id="container" class="mol-container"></div>
+        <script>
+            let pdb = `{mol}`; // Use template literal to properly escape PDB content
+            $(document).ready(function () {{
+                let element = $("#container");
+                let config = {{ backgroundColor: "white" }};
+                let viewer = $3Dmol.createViewer(element, config);
+                
+                {high_score_script}
+                
+                // Add hover functionality
+                viewer.setHoverable(
+                    {{}}, 
+                    true, 
+                    function(atom, viewer, event, container) {{
+                        if (!atom.label) {{
+                            atom.label = viewer.addLabel(
+                                atom.resn + ":" +atom.resi + ":" + atom.atom, 
+                                {{
+                                    position: atom, 
+                                    backgroundColor: 'mintcream', 
+                                    fontColor: 'black',
+                                    fontSize: 18,
+                                    padding: 4
+                                }}
+                            );
+                        }}
+                    }},
+                    function(atom, viewer) {{
+                        if (atom.label) {{
+                            viewer.removeLabel(atom.label);
+                            delete atom.label;
+                        }}
+                    }}
+                );
+                
+                viewer.zoomTo();
+                viewer.render();
+                viewer.zoom(0.8, 2000);
+            }});
+        </script>
+    </body>
+    </html>
+    """
+    
+    # Return the HTML content within an iframe safely encoded for special characters
+    return f'<iframe width="100%" height="700" srcdoc="{html_content.replace(chr(34), "&quot;").replace(chr(39), "&#39;")}"></iframe>'
+
+with gr.Blocks(css="""
+    /* Customize Gradio button colors */
+    #visualize-btn, #predict-btn {
+        background-color: #FF7300; /* Deep orange */
+        color: white;
+        border-radius: 5px;
+        padding: 10px;
+        font-weight: bold;
+    }
+    #visualize-btn:hover, #predict-btn:hover {
+        background-color: #CC5C00; /* Darkened orange on hover */
+    }
+""") as demo:
+    gr.Markdown("# Protein Binding Site Prediction")
+    
+    # Mode selection
+    mode = gr.Radio(
+        choices=["PDB ID", "Upload File"],
+        value="PDB ID",
+        label="Input Mode",
+        info="Choose whether to input a PDB ID or upload a PDB/CIF file."
+    )
+
+    # Input components based on mode
+    pdb_input = gr.Textbox(value="2F6V", label="PDB ID", placeholder="Enter PDB ID here...")
+    pdb_file = gr.File(label="Upload PDB/CIF File", visible=False)
+    visualize_btn = gr.Button("Visualize Structure", elem_id="visualize-btn")
+
+    molecule_output2 = Molecule3D(label="Protein Structure", reps=[
+        {
+            "model": 0,
+            "style": "cartoon",
+            "color": "whiteCarbon",
+            "residue_range": "",
+            "around": 0,
+            "byres": False,
+        }
+    ])
+
+    with gr.Row():
+        segment_input = gr.Textbox(value="A", label="Chain ID (protein)", placeholder="Enter Chain ID here...",
+        info="Choose in which chain to predict binding sites.")
+        prediction_btn = gr.Button("Predict Binding Site", elem_id="predict-btn")
+
+    # Add score type selector
+    score_type = gr.Radio(
+        choices=["Normalized Scores", "Raw Scores"],
+        value="Normalized Scores",
+        label="Score Visualization Type",
+        info="Choose which score type to visualize"
+    )
+
+    molecule_output = gr.HTML(label="Protein Structure")
+    explanation_vis = gr.Markdown("""
+    Score dependent colorcoding:
+    - 0.0-0.2: white  
+    - 0.2–0.4: light orange  
+    - 0.4–0.6: yellow orange
+    - 0.6–0.8: orange
+    - 0.8–1.0: red
+    """)
+    predictions_output = gr.Textbox(label="Visualize Prediction with PyMol")
+    gr.Markdown("### Download:\n- List of predicted binding site residues\n- PDB with score in beta factor column")
+    download_output = gr.File(label="Download Files", file_count="multiple")
+    
+    # Store these as state variables so we can switch between them
+    raw_scores_state = gr.State(None)
+    norm_scores_state = gr.State(None)
+    last_pdb_path = gr.State(None)
+    last_segment = gr.State(None)
+    last_pdb_id = gr.State(None)
+    
+    def process_interface(mode, pdb_id, pdb_file, chain_id, score_type_val):
+        selected_score_type = 'normalized' if score_type_val == "Normalized Scores" else 'raw'
+        
+        # First get the actual PDB file path
+        if mode == "PDB ID":
+            pdb_path = fetch_pdb(pdb_id)  # Get the actual file path
+            
+            pymol_cmd, mol_vis, files, raw_scores, norm_scores, pdb_id_result, segment = process_pdb(pdb_path, chain_id, selected_score_type)
+            # Store the actual file path, not just the PDB ID
+            return pymol_cmd, mol_vis, files, raw_scores, norm_scores, pdb_path, chain_id, pdb_id_result
+        elif mode == "Upload File":
+            _, ext = os.path.splitext(pdb_file.name)
+            file_path = os.path.join('./', f"{_}{ext}")
+            if ext == '.cif':
+                pdb_path = convert_cif_to_pdb(file_path)
+            else:
+                pdb_path = file_path
+            
+            pymol_cmd, mol_vis, files, raw_scores, norm_scores, pdb_id_result, segment = process_pdb(pdb_path, chain_id, selected_score_type)
+            return pymol_cmd, mol_vis, files, raw_scores, norm_scores, pdb_path, chain_id, pdb_id_result
+    
+    def update_visualization_and_files(score_type_val, raw_scores, norm_scores, pdb_path, segment, pdb_id):
+        if raw_scores is None or norm_scores is None or pdb_path is None or segment is None or pdb_id is None:
+            return None, None, None
+        
+        # Choose scores based on radio button selection
+        selected_score_type = 'normalized' if score_type_val == "Normalized Scores" else 'raw'
+        selected_scores = norm_scores if selected_score_type == 'normalized' else raw_scores
+        
+        # Generate visualization with selected scores
+        mol_vis = molecule(pdb_path, selected_scores, segment)
+        
+        # Generate PyMOL commands and downloadable files
+        # Get structure for residue info
+        _, ext = os.path.splitext(pdb_path)
+        parser = MMCIFParser(QUIET=True) if ext == '.cif' else PDBParser(QUIET=True)
+        structure = parser.get_structure('protein', pdb_path)
+        chain = structure[0][segment]
+        protein_residues = [res for res in chain if is_aa(res)]
+        sequence = "".join(seq1(res.resname) for res in protein_residues)
+        
+        # Define score brackets
+        score_brackets = {
+            "0.0-0.2": (0.0, 0.2),
+            "0.2-0.4": (0.2, 0.4),
+            "0.4-0.6": (0.4, 0.6),
+            "0.6-0.8": (0.6, 0.8),
+            "0.8-1.0": (0.8, 1.0)
+        }
+        
+        # Initialize a dictionary to store residues by bracket
+        residues_by_bracket = {bracket: [] for bracket in score_brackets}
+        
+        # Categorize residues into brackets
+        for resi, score in selected_scores:
+            for bracket, (lower, upper) in score_brackets.items():
+                if lower <= score < upper:
+                    residues_by_bracket[bracket].append(resi)
+                    break
+        
+        # Generate timestamp
+        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        
+        # Generate result text and PyMOL commands based on score type
+        display_score_type = "Normalized" if selected_score_type == 'normalized' else "Raw"
+        scores_array = [score for _, score in selected_scores]
+        result_str = generate_results_text(pdb_id, segment, residues_by_bracket, protein_residues, sequence, 
+                                           scores_array, current_time, display_score_type)
+        pymol_commands = generate_pymol_commands(pdb_id, segment, residues_by_bracket, current_time, display_score_type)
+        
+        # Create chain-specific PDB with scores in B-factor
+        scored_pdb = create_chain_specific_pdb(pdb_path, segment, selected_scores, protein_residues)
+        
+        # Create prediction file
+        prediction_file = f"{pdb_id}_{display_score_type.lower()}_binding_site_residues.txt"
+        with open(prediction_file, "w") as f:
+            f.write(result_str)
+        
+        scored_pdb_name = f"{pdb_id}_{segment}_{display_score_type.lower()}_predictions_scores.pdb"
+        os.rename(scored_pdb, scored_pdb_name)
+        
+        return mol_vis, pymol_commands, [prediction_file, scored_pdb_name]
+
+    def fetch_interface(mode, pdb_id, pdb_file):
+        if mode == "PDB ID":
+            return fetch_pdb(pdb_id)
+        elif mode == "Upload File":
+            _, ext = os.path.splitext(pdb_file.name)
+            file_path = os.path.join('./', f"{_}{ext}")
+            if ext == '.cif':
+                pdb_path = convert_cif_to_pdb(file_path)
+            else:
+                pdb_path= file_path
+            return pdb_path
+
+    def toggle_mode(selected_mode):
+        if selected_mode == "PDB ID":
+            return gr.update(visible=True), gr.update(visible=False)
+        else:
+            return gr.update(visible=False), gr.update(visible=True)
+
+    
+        
+    mode.change(
+        toggle_mode,
+        inputs=[mode],
+        outputs=[pdb_input, pdb_file]
+    )
+
+    prediction_btn.click(
+        process_interface, 
+        inputs=[mode, pdb_input, pdb_file, segment_input, score_type], 
+        outputs=[predictions_output, molecule_output, download_output, 
+                raw_scores_state, norm_scores_state, last_pdb_path, last_segment, last_pdb_id]
+    )
+    
+    # Update visualization, PyMOL commands, and files when score type changes
+    score_type.change(
+        update_visualization_and_files,
+        inputs=[score_type, raw_scores_state, norm_scores_state, last_pdb_path, last_segment, last_pdb_id],
+        outputs=[molecule_output, predictions_output, download_output]
+    )
+
+    visualize_btn.click(
+        fetch_interface, 
+        inputs=[mode, pdb_input, pdb_file], 
+        outputs=molecule_output2
+    )
+
+    gr.Markdown("## Examples")
+    gr.Examples(
+        examples=[
+            ["7RPZ", "A"],
+            ["2IWI", "B"],
+            ["7LCJ", "R"],
+            ["4OBE", "A"]
+        ],
+        inputs=[pdb_input, segment_input],
+        outputs=[predictions_output, molecule_output, download_output]
+    )
+
+    def predict_utils(sequence):
+        input_ids = tokenizer(" ".join(sequence), return_tensors="pt").input_ids.to(device)
+        with torch.no_grad():
+            outputs = model(input_ids).logits.detach().cpu().numpy().squeeze()
+    
+        raw_scores = expit(outputs[:, 1] - outputs[:, 0])
+        normalized_scores = normalize_scores(raw_scores)
+    
+        return {
+        "raw_scores": raw_scores.tolist(),
+        "normalized_scores": normalized_scores.tolist()
+    }
+
+    dummy_input = gr.Textbox(visible=False)
+    dummy_output = gr.Textbox(visible=False)
+    
+    dummy_btn = gr.Button("Predict Sequence", visible=False)
+    dummy_btn.click(
+        predict_utils,
+        inputs=[dummy_input],
+        outputs=[dummy_output]
+    )
+
+demo.launch(share=True)

model_loader.py

@@ -0,0 +1,640 @@
+from huggingface_hub import hf_hub_download
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from torch.utils.data import DataLoader
+
+import re
+import numpy as np
+import os
+import pandas as pd
+import copy
+
+import transformers, datasets
+from transformers.modeling_outputs import TokenClassifierOutput
+from transformers.models.t5.modeling_t5 import T5Config, T5PreTrainedModel, T5Stack
+from transformers.utils.model_parallel_utils import assert_device_map, get_device_map
+from transformers import T5EncoderModel, T5Tokenizer
+from transformers.models.esm.modeling_esm import EsmPreTrainedModel, EsmModel
+from transformers import AutoTokenizer
+from transformers import TrainingArguments, Trainer, set_seed
+from transformers import DataCollatorForTokenClassification
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Union
+
+# for custom DataCollator
+from transformers.data.data_collator import DataCollatorMixin
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+from transformers.utils import PaddingStrategy
+
+from datasets import Dataset
+
+from scipy.special import expit
+
+#import peft
+#from peft import get_peft_config, PeftModel, PeftConfig, inject_adapter_in_model, LoraConfig
+
+cnn_head=True #False set True for Rostlab/prot_t5_xl_half_uniref50-enc
+ffn_head=False #False
+transformer_head=False
+custom_lora=True #False #only true for Rostlab/prot_t5_xl_half_uniref50-enc
+
+class ClassConfig:
+    def __init__(self, dropout=0.2, num_labels=3):
+        self.dropout_rate = dropout
+        self.num_labels = num_labels
+
+class T5EncoderForTokenClassification(T5PreTrainedModel):
+
+    def __init__(self, config: T5Config, class_config: ClassConfig):
+        super().__init__(config)
+        self.num_labels = class_config.num_labels
+        self.config = config
+
+        self.shared = nn.Embedding(config.vocab_size, config.d_model)
+
+        encoder_config = copy.deepcopy(config)
+        encoder_config.use_cache = False
+        encoder_config.is_encoder_decoder = False
+        self.encoder = T5Stack(encoder_config, self.shared)
+
+        self.dropout = nn.Dropout(class_config.dropout_rate)
+
+        # Initialize different heads based on class_config
+        if cnn_head:
+            self.cnn = nn.Conv1d(config.hidden_size, 512, kernel_size=3, padding=1)
+            self.classifier = nn.Linear(512, class_config.num_labels)
+        elif ffn_head:
+            # Multi-layer feed-forward network (FFN) head
+            self.ffn = nn.Sequential(
+                nn.Linear(config.hidden_size, 512),
+                nn.ReLU(),
+                nn.Linear(512, 256),
+                nn.ReLU(),
+                nn.Linear(256, class_config.num_labels)
+            )
+        elif transformer_head:
+            # Transformer layer head
+            encoder_layer = nn.TransformerEncoderLayer(d_model=config.hidden_size, nhead=8)
+            self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=1)
+            self.classifier = nn.Linear(config.hidden_size, class_config.num_labels)
+        else:
+            # Default classification head
+            self.classifier = nn.Linear(config.hidden_size, class_config.num_labels)
+        
+        self.post_init()
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+    def parallelize(self, device_map=None):
+        self.device_map = (
+            get_device_map(len(self.encoder.block), range(torch.cuda.device_count()))
+            if device_map is None
+            else device_map
+        )
+        assert_device_map(self.device_map, len(self.encoder.block))
+        self.encoder.parallelize(self.device_map)
+        self.classifier = self.classifier.to(self.encoder.first_device)
+        self.model_parallel = True
+
+    def deparallelize(self):
+        self.encoder.deparallelize()
+        self.encoder = self.encoder.to("cpu")
+        self.model_parallel = False
+        self.device_map = None
+        torch.cuda.empty_cache()
+
+    def get_input_embeddings(self):
+        return self.shared
+
+    def set_input_embeddings(self, new_embeddings):
+        self.shared = new_embeddings
+        self.encoder.set_input_embeddings(new_embeddings)
+
+    def get_encoder(self):
+        return self.encoder
+
+    def _prune_heads(self, heads_to_prune):
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            head_mask=head_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = outputs[0]
+        sequence_output = self.dropout(sequence_output)
+
+        # Forward pass through the selected head
+        if cnn_head:
+            # CNN head
+            sequence_output = sequence_output.permute(0, 2, 1)  # Prepare shape for CNN
+            cnn_output = self.cnn(sequence_output)
+            cnn_output = F.relu(cnn_output)
+            cnn_output = cnn_output.permute(0, 2, 1)  # Shape back for classifier
+            logits = self.classifier(cnn_output)
+        elif ffn_head:
+            # FFN head
+            logits = self.ffn(sequence_output)
+        elif transformer_head:
+            # Transformer head
+            transformer_output = self.transformer_encoder(sequence_output)
+            logits = self.classifier(transformer_output)
+        else:
+            # Default classification head
+            logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            active_loss = attention_mask.view(-1) == 1
+            active_logits = logits.view(-1, self.num_labels)
+            active_labels = torch.where(
+                active_loss, labels.view(-1), torch.tensor(-100).type_as(labels)
+            )
+            valid_logits = active_logits[active_labels != -100]
+            valid_labels = active_labels[active_labels != -100]
+            valid_labels = valid_labels.to(valid_logits.device)
+            valid_labels = valid_labels.long()
+            loss = loss_fct(valid_logits, valid_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+# Modifies an existing transformer and introduce the LoRA layers
+
+class CustomLoRAConfig:
+    def __init__(self):
+        self.lora_rank = 4
+        self.lora_init_scale = 0.01
+        self.lora_modules = ".*SelfAttention|.*EncDecAttention"
+        self.lora_layers = "q|k|v|o"
+        self.trainable_param_names = ".*layer_norm.*|.*lora_[ab].*"
+        self.lora_scaling_rank = 1
+        # lora_modules and lora_layers are speicified with regular expressions
+        # see https://www.w3schools.com/python/python_regex.asp for reference
+        
+class LoRALinear(nn.Module):
+    def __init__(self, linear_layer, rank, scaling_rank, init_scale):
+        super().__init__()
+        self.in_features = linear_layer.in_features
+        self.out_features = linear_layer.out_features
+        self.rank = rank
+        self.scaling_rank = scaling_rank
+        self.weight = linear_layer.weight
+        self.bias = linear_layer.bias
+        if self.rank > 0:
+            self.lora_a = nn.Parameter(torch.randn(rank, linear_layer.in_features) * init_scale)
+            if init_scale < 0:
+                self.lora_b = nn.Parameter(torch.randn(linear_layer.out_features, rank) * init_scale)
+            else:
+                self.lora_b = nn.Parameter(torch.zeros(linear_layer.out_features, rank))
+        if self.scaling_rank:
+            self.multi_lora_a = nn.Parameter(
+                torch.ones(self.scaling_rank, linear_layer.in_features)
+                + torch.randn(self.scaling_rank, linear_layer.in_features) * init_scale
+            )
+            if init_scale < 0:
+                self.multi_lora_b = nn.Parameter(
+                    torch.ones(linear_layer.out_features, self.scaling_rank)
+                    + torch.randn(linear_layer.out_features, self.scaling_rank) * init_scale
+                )
+            else:
+                self.multi_lora_b = nn.Parameter(torch.ones(linear_layer.out_features, self.scaling_rank))
+
+    def forward(self, input):
+        if self.scaling_rank == 1 and self.rank == 0:
+            # parsimonious implementation for ia3 and lora scaling
+            if self.multi_lora_a.requires_grad:
+                hidden = F.linear((input * self.multi_lora_a.flatten()), self.weight, self.bias)
+            else:
+                hidden = F.linear(input, self.weight, self.bias)
+            if self.multi_lora_b.requires_grad:
+                hidden = hidden * self.multi_lora_b.flatten()
+            return hidden
+        else:
+            # general implementation for lora (adding and scaling)
+            weight = self.weight
+            if self.scaling_rank:
+                weight = weight * torch.matmul(self.multi_lora_b, self.multi_lora_a) / self.scaling_rank
+            if self.rank:
+                weight = weight + torch.matmul(self.lora_b, self.lora_a) / self.rank
+            return F.linear(input, weight, self.bias)
+
+    def extra_repr(self):
+        return "in_features={}, out_features={}, bias={}, rank={}, scaling_rank={}".format(
+            self.in_features, self.out_features, self.bias is not None, self.rank, self.scaling_rank
+        )
+
+
+def modify_with_lora(transformer, config):
+    for m_name, module in dict(transformer.named_modules()).items():
+        if re.fullmatch(config.lora_modules, m_name):
+            for c_name, layer in dict(module.named_children()).items():
+                if re.fullmatch(config.lora_layers, c_name):
+                    assert isinstance(
+                        layer, nn.Linear
+                    ), f"LoRA can only be applied to torch.nn.Linear, but {layer} is {type(layer)}."
+                    setattr(
+                        module,
+                        c_name,
+                        LoRALinear(layer, config.lora_rank, config.lora_scaling_rank, config.lora_init_scale),
+                    )
+    return transformer
+
+
+def load_T5_model_classification(checkpoint, num_labels, half_precision, full = False, deepspeed=True):
+    # Load model and tokenizer
+
+    if "ankh" in checkpoint :
+        model = T5EncoderModel.from_pretrained(checkpoint,resume_download=True)
+        tokenizer = AutoTokenizer.from_pretrained(checkpoint,resume_download=True)
+
+    elif "prot_t5" in checkpoint:
+        # possible to load the half precision model (thanks to @pawel-rezo for pointing that out)
+        if half_precision and deepspeed:
+            #tokenizer = T5Tokenizer.from_pretrained('Rostlab/prot_t5_xl_half_uniref50-enc', do_lower_case=False)
+            #model = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_half_uniref50-enc", torch_dtype=torch.float16)#.to(torch.device('cuda')
+            tokenizer = T5Tokenizer.from_pretrained(checkpoint, do_lower_case=False,resume_download=True)
+            model = T5EncoderModel.from_pretrained(checkpoint, torch_dtype=torch.float16).to(torch.device('cuda'),resume_download=True)
+        else:
+            model = T5EncoderModel.from_pretrained(checkpoint)
+            tokenizer = T5Tokenizer.from_pretrained(checkpoint)
+                
+    elif "ProstT5" in checkpoint:
+        if half_precision and deepspeed: 
+            tokenizer = T5Tokenizer.from_pretrained(checkpoint, do_lower_case=False,resume_download=True)
+            model = T5EncoderModel.from_pretrained(checkpoint, torch_dtype=torch.float16).to(torch.device('cuda'),resume_download=True)
+        else:
+            model = T5EncoderModel.from_pretrained(checkpoint,resume_download=True)
+            tokenizer = T5Tokenizer.from_pretrained(checkpoint,resume_download=True) 
+    
+    # Create new Classifier model with PT5 dimensions
+    class_config=ClassConfig(num_labels=num_labels)
+    class_model=T5EncoderForTokenClassification(model.config,class_config)
+    
+    # Set encoder and embedding weights to checkpoint weights
+    class_model.shared=model.shared
+    class_model.encoder=model.encoder    
+    
+    # Delete the checkpoint model
+    model=class_model
+    del class_model
+    
+    if full == True:
+        return model, tokenizer 
+    
+    # Print number of trainable parameters
+    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+    params = sum([np.prod(p.size()) for p in model_parameters])
+    print("T5_Classfier\nTrainable Parameter: "+ str(params))    
+
+    if custom_lora:
+        #the linear CustomLoRAConfig allows better quality predictions, but more memory is needed
+        # Add model modification lora
+        config = CustomLoRAConfig()
+        
+        # Add LoRA layers
+        model = modify_with_lora(model, config)
+        
+        # Freeze Embeddings and Encoder (except LoRA)
+        for (param_name, param) in model.shared.named_parameters():
+                    param.requires_grad = False
+        for (param_name, param) in model.encoder.named_parameters():
+                    param.requires_grad = False       
+    
+        for (param_name, param) in model.named_parameters():
+                if re.fullmatch(config.trainable_param_names, param_name):
+                    param.requires_grad = True
+
+    else:
+        # lora modification
+        peft_config = LoraConfig(
+            r=4, lora_alpha=1, bias="all", target_modules=["q","k","v","o"]
+        )
+        
+        model = inject_adapter_in_model(peft_config, model)
+        
+        # Unfreeze the prediction head
+        for (param_name, param) in model.classifier.named_parameters():
+                    param.requires_grad = True  
+
+    # Print trainable Parameter          
+    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+    params = sum([np.prod(p.size()) for p in model_parameters])
+    print("T5_LoRA_Classfier\nTrainable Parameter: "+ str(params) + "\n")
+    
+    return model, tokenizer
+
+class EsmForTokenClassificationCustom(EsmPreTrainedModel):
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+    _keys_to_ignore_on_load_missing = [r"position_ids", r"cnn", r"ffn", r"transformer"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.esm = EsmModel(config, add_pooling_layer=False)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        if cnn_head:
+            self.cnn = nn.Conv1d(config.hidden_size, 512, kernel_size=3, padding=1)
+            self.classifier = nn.Linear(512, config.num_labels)
+        elif ffn_head:
+            # Multi-layer feed-forward network (FFN) as an alternative head
+            self.ffn = nn.Sequential(
+                nn.Linear(config.hidden_size, 512),
+                nn.ReLU(),
+                nn.Linear(512, 256),
+                nn.ReLU(),
+                nn.Linear(256, config.num_labels)
+            )
+        elif transformer_head:
+            # Transformer layer as an alternative head
+            encoder_layer = nn.TransformerEncoderLayer(d_model=config.hidden_size, nhead=8)
+            self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=1)
+            self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        else:
+            # Default classification head
+            self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        self.init_weights()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, TokenClassifierOutput]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        outputs = self.esm(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        
+        sequence_output = outputs[0]
+        sequence_output = self.dropout(sequence_output)
+
+        if cnn_head:
+            sequence_output = sequence_output.transpose(1, 2)
+            sequence_output = self.cnn(sequence_output)
+            sequence_output = sequence_output.transpose(1, 2)
+            logits = self.classifier(sequence_output)
+        elif ffn_head:
+            logits = self.ffn(sequence_output)
+        elif transformer_head:
+            # Apply transformer encoder for the transformer head
+            sequence_output = self.transformer_encoder(sequence_output)
+            logits = self.classifier(sequence_output)
+        else:
+            logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            active_loss = attention_mask.view(-1) == 1
+            active_logits = logits.view(-1, self.num_labels)
+            active_labels = torch.where(
+                active_loss, labels.view(-1), torch.tensor(-100).type_as(labels)
+            )
+            valid_logits = active_logits[active_labels != -100]
+            valid_labels = active_labels[active_labels != -100]
+            valid_labels = valid_labels.type(torch.LongTensor).to('cuda:0')
+            loss = loss_fct(valid_logits, valid_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear) or isinstance(module, nn.Conv1d):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+# based on transformers DataCollatorForTokenClassification
+@dataclass
+class DataCollatorForTokenClassificationESM(DataCollatorMixin):
+    """
+    Data collator that will dynamically pad the inputs received, as well as the labels.
+    Args:
+        tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):
+            The tokenizer used for encoding the data.
+        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):
+            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
+            among:
+            - `True` or `'longest'` (default): Pad to the longest sequence in the batch (or no padding if only a single
+              sequence is provided).
+            - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
+              acceptable input length for the model if that argument is not provided.
+            - `False` or `'do_not_pad'`: No padding (i.e., can output a batch with sequences of different lengths).
+        max_length (`int`, *optional*):
+            Maximum length of the returned list and optionally padding length (see above).
+        pad_to_multiple_of (`int`, *optional*):
+            If set will pad the sequence to a multiple of the provided value.
+            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
+            7.5 (Volta).
+        label_pad_token_id (`int`, *optional*, defaults to -100):
+            The id to use when padding the labels (-100 will be automatically ignore by PyTorch loss functions).
+        return_tensors (`str`):
+            The type of Tensor to return. Allowable values are "np", "pt" and "tf".
+    """
+
+    tokenizer: PreTrainedTokenizerBase
+    padding: Union[bool, str, PaddingStrategy] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+    label_pad_token_id: int = -100
+    return_tensors: str = "pt"
+
+    def torch_call(self, features):
+        import torch
+
+        label_name = "label" if "label" in features[0].keys() else "labels"
+        labels = [feature[label_name] for feature in features] if label_name in features[0].keys() else None
+
+        no_labels_features = [{k: v for k, v in feature.items() if k != label_name} for feature in features]
+
+        batch = self.tokenizer.pad(
+            no_labels_features,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors="pt",
+        )
+
+        if labels is None:
+            return batch
+
+        sequence_length = batch["input_ids"].shape[1]
+        padding_side = self.tokenizer.padding_side
+
+        def to_list(tensor_or_iterable):
+            if isinstance(tensor_or_iterable, torch.Tensor):
+                return tensor_or_iterable.tolist()
+            return list(tensor_or_iterable)
+
+        if padding_side == "right":
+            batch[label_name] = [
+                # to_list(label) + [self.label_pad_token_id] * (sequence_length - len(label)) for label in labels
+                # changed to pad the special tokens at the beginning and end of the sequence
+                [self.label_pad_token_id] + to_list(label) + [self.label_pad_token_id] * (sequence_length - len(label)-1) for label in labels
+            ]
+        else:
+            batch[label_name] = [
+                [self.label_pad_token_id] * (sequence_length - len(label)) + to_list(label) for label in labels
+            ]
+
+        batch[label_name] = torch.tensor(batch[label_name], dtype=torch.float)
+        return batch
+
+def _torch_collate_batch(examples, tokenizer, pad_to_multiple_of: Optional[int] = None):
+    """Collate `examples` into a batch, using the information in `tokenizer` for padding if necessary."""
+    import torch
+
+    # Tensorize if necessary.
+    if isinstance(examples[0], (list, tuple, np.ndarray)):
+        examples = [torch.tensor(e, dtype=torch.long) for e in examples]
+
+    length_of_first = examples[0].size(0)
+
+    # Check if padding is necessary.
+
+    are_tensors_same_length = all(x.size(0) == length_of_first for x in examples)
+    if are_tensors_same_length and (pad_to_multiple_of is None or length_of_first % pad_to_multiple_of == 0):
+        return torch.stack(examples, dim=0)
+
+    # If yes, check if we have a `pad_token`.
+    if tokenizer._pad_token is None:
+        raise ValueError(
+            "You are attempting to pad samples but the tokenizer you are using"
+            f" ({tokenizer.__class__.__name__}) does not have a pad token."
+        )
+
+    # Creating the full tensor and filling it with our data.
+    max_length = max(x.size(0) for x in examples)
+    if pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
+        max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of
+    result = examples[0].new_full([len(examples), max_length], tokenizer.pad_token_id)
+    for i, example in enumerate(examples):
+        if tokenizer.padding_side == "right":
+            result[i, : example.shape[0]] = example
+        else:
+            result[i, -example.shape[0] :] = example
+    return result
+
+def tolist(x):
+    if isinstance(x, list):
+        return x
+    elif hasattr(x, "numpy"):  # Checks for TF tensors without needing the import
+        x = x.numpy()
+    return x.tolist()
+
+#load ESM2 models
+def load_esm_model_classification(checkpoint, num_labels, half_precision, full=False, deepspeed=True):
+    
+    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+
+    
+    if half_precision and deepspeed:
+        model = EsmForTokenClassificationCustom.from_pretrained(checkpoint, 
+                                                                num_labels = num_labels, 
+                                                                ignore_mismatched_sizes=True,
+                                                                torch_dtype = torch.float16)
+    else:
+        model = EsmForTokenClassificationCustom.from_pretrained(checkpoint, 
+                                                                num_labels = num_labels,
+                                                                ignore_mismatched_sizes=True)
+        
+    if full == True:
+        return model, tokenizer 
+        
+    peft_config = LoraConfig(
+        r=4, lora_alpha=1, bias="all", target_modules=["query","key","value","dense"]
+    )
+    
+    model = inject_adapter_in_model(peft_config, model)
+
+    #model.gradient_checkpointing_enable()
+    
+    # Unfreeze the prediction head
+    for (param_name, param) in model.classifier.named_parameters():
+                param.requires_grad = True  
+    
+    return model, tokenizer
+
+def load_model(checkpoint,max_length):
+    #checkpoint='ThorbenF/prot_t5_xl_uniref50'
+    #best_model_path='ThorbenF/prot_t5_xl_uniref50/cpt.pth'
+    full=False
+    deepspeed=False
+    mixed=False 
+    num_labels=2
+    
+    print(checkpoint, num_labels, mixed, full, deepspeed)
+            
+    # Determine model type and load accordingly
+    if "esm" in checkpoint:
+        model, tokenizer = load_esm_model_classification(checkpoint, num_labels, mixed, full, deepspeed)
+    else:
+        model, tokenizer = load_T5_model_classification(checkpoint, num_labels, mixed, full, deepspeed)
+    
+
+    # Download the file
+    local_file = hf_hub_download(repo_id=checkpoint, filename="cpt.pth")
+
+    # Load the best model state
+    state_dict = torch.load(local_file, map_location=torch.device('cpu'), weights_only=True)
+    model.load_state_dict(state_dict)
+
+    return model, tokenizer

requirements.txt

@@ -0,0 +1,14 @@
+torch>=1.13.0
+transformers>=4.30.0
+datasets>=2.9.0
+peft>=0.0.7
+scipy>=1.7.0
+pandas>=1.1.0
+numpy>=1.19.0
+scikit-learn>=0.24.0
+sentencepiece
+huggingface_hub>=0.15.0
+requests
+gradio_molecule3d
+biopython>=1.81
+pydantic==2.1.1