Spaces:

Esperanto
/

Protein-Llama-3-8B-Gradio

Runtime error

File size: 10,472 Bytes

import os
import torch
import transformers
from transformers import GenerationConfig, pipeline, AutoTokenizer, AutoModelForCausalLM, EsmForProteinFolding
import os
import tempfile
import subprocess
import pandas as pd
import numpy as np
import gradio as gr
from time import time
import requests


model_id = "Esperanto/Protein-Phi-3-mini"

model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=torch.float16,
        trust_remote_code=True,
    )


tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct", trust_remote_code=True)

tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"

#Creating the pipeline for generation
generator = pipeline('text-generation', model=model, tokenizer=tokenizer)


# Loading the ESM Model
esm_model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1")
esm_tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")


#Ensures that final output contains only valid amino acids
def clean_protein_sequence(protein_seq):
    # Valid amino acid characters
    valid_amino_acids = "ACDEFGHIKLMNPQRSTVWY"
    
    # Filter out any characters that are not valid amino acids
    cleaned_seq = ''.join([char for char in protein_seq if char in valid_amino_acids])
    
    return cleaned_seq

#convert pLDDT to percentage
def modify_b_factors(pdb_content, multiplier):
    modified_pdb = []
    for line in pdb_content.split('\n'):
        if line.startswith("ATOM"):
            b_factor = float(line[60:66].strip())
            new_b_factor = b_factor * multiplier
            new_line = f"{line[:60]}{new_b_factor:6.2f}{line[66:]}"
            modified_pdb.append(new_line)
        else:
            modified_pdb.append(line)
    return "\n".join(modified_pdb)

#saves the structure output from ESMFold as a PDB file in a temporary folder
def save_pdb(input_sequence):
    inputs = esm_tokenizer([input_sequence], return_tensors="pt", add_special_tokens=False)
    with torch.no_grad():
        outputs = esm_model(**inputs)
    pdb_string_unscaled = esm_model.output_to_pdb(outputs)[0]
    pdb_string = modify_b_factors(pdb_string_unscaled, 100)
    plddt_values = outputs.plddt.tolist()[0][0]
    plddt_values = [round(value * 100, 2) for value in plddt_values]
    file_path = os.path.join('Protein-Llama-3-8B-Gradio/temporary_folder', f"protein.pdb")
    os.makedirs(os.path.dirname(file_path), exist_ok=True)
    with open(file_path, "w") as f:
        f.write(pdb_string)
    
    return np.mean(plddt_values)

#reads the PDB file
def read_prot(molpath):
    with open(molpath, "r") as fp:
        lines = fp.readlines()
    mol = ""
    for l in lines:
        mol += l
    return mol

def get_cov2_pdb():
    pdb_id = '6vxx'
    url = f'https://files.rcsb.org/download/{pdb_id}.pdb'

    # Fetch the PDB file from the RCSB PDB website
    response = requests.get(url)

    # Check if the request was successful
    mol = response.text

    return mol


def protein_visual_html(input_pdb):

    if input_pdb == 'cov2':
        mol = get_cov2_pdb()

    else:
        mol = read_prot(input_pdb)

    x = (
        """<!DOCTYPE html>
        <html>
        <head>    
    <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
    <style>
    body{
        font-family:sans-serif
    }
    .mol-container {
    width: 100%;
    height: 600px;
    position: relative;
    }
    .mol-container select{
        background-image:None;
    }
    </style>
     <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js" integrity="sha512-STof4xm1wgkfm7heWqFJVn58Hm3EtS31XFaagaa8VMReCXAkQnJZ+jEy8PCC/iT18dFy95WcExNHFTqLyp72eQ==" crossorigin="anonymous" referrerpolicy="no-referrer"></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 + """`  
      
             $(document).ready(function () {
                let element = $("#container");
                let config = { backgroundColor: "white" };
                let viewer = $3Dmol.createViewer(element, config);
                viewer.addModel(pdb, "pdb");
                viewer.getModel(0).setStyle({}, { cartoon: { color:"spectrum" } });
                viewer.zoomTo();
                viewer.render();
                viewer.zoom(0.8, 2000);
              })
        </script>
        </body></html>"""
    )

    return f"""<iframe style="width: 100%; height: 600px" name="result" allow="midi; geolocation; microphone; camera; 
    display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
    allow-scripts allow-same-origin allow-popups 
    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
    allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""


def predict_structure(input_sequence):
    #Hard coding the SARS-CoV 2 protein sequence and structure for instant demo purposes
    if input_sequence == 'SNASADAQSFLNRVCGVSAARLTPCGTGTSTDVVYRAFDIYNDKVAGFAKFLKTNCCRFQEKDEDDNLIDSYFVVKRHTFSNYQHEETIYNLLKDCPAVAKHDFFKFRIDGDMVPHISRQRLTKYTMADLVYALRHFDEGNCDTLKEILVTYNCCDDDYFNKKDWYDFVENPDILRVYANLGERVRQALLKTVQFCDAMRNAGIVGVLTLDNQDLNGNWYDFGDFIQTTPGSGVPVVDSYYSLLMPILTLTRALTAESHVDTDLTKPYIKWDLLKYDFTEERLKLFDRYFKYWDQTYHPNCVNCLDDRCILHCANFNVLFSTVFPPTSFGPLVRKIFVDGVPFVVSTGYHFRELGVVHNQDVNLHSSRLSFKELLVYAADPAMHAASGNLLLDKRTTCFSVAALTNNVAFQTVKPGNFNKDFYDFAVSKGFFKEGSSVELKHFFFAQDGNAAISDYDYYRYNLPTMCDIRQLLFVVEVVDKYFDCYDGGCINANQVI':
        return protein_visual_html('Protein-Llama-3-8B-Gradio/sars_cov.pdb')
    else:
        plddt = save_pdb(input_sequence)
        #Creating HTML visualization for the PDB file stores in temporary folder
        pdb_path = 'cov2'
        return protein_visual_html(pdb_path)

def generate_protein_sequence(sequence, seq_length, property=''):
    enzymes = ["Non-Hemolytic", "Soluble", "Oxidoreductase", "Transferase", "Hydrolase", "Lyase", "Isomerase", "Ligase", "Translocase"]
    start_time = time()
    
    if property is None:
        input_prompt = 'Seq=<' + sequence
    elif property == 'SARS-CoV-2 Spike Protein (example)':
        cleaned_seq = 'SNASADAQSFLNRVCGVSAARLTPCGTGTSTDVVYRAFDIYNDKVAGFAKFLKTNCCRFQEKDEDDNLIDSYFVVKRHTFSNYQHEETIYNLLKDCPAVAKHDFFKFRIDGDMVPHISRQRLTKYTMADLVYALRHFDEGNCDTLKEILVTYNCCDDDYFNKKDWYDFVENPDILRVYANLGERVRQALLKTVQFCDAMRNAGIVGVLTLDNQDLNGNWYDFGDFIQTTPGSGVPVVDSYYSLLMPILTLTRALTAESHVDTDLTKPYIKWDLLKYDFTEERLKLFDRYFKYWDQTYHPNCVNCLDDRCILHCANFNVLFSTVFPPTSFGPLVRKIFVDGVPFVVSTGYHFRELGVVHNQDVNLHSSRLSFKELLVYAADPAMHAASGNLLLDKRTTCFSVAALTNNVAFQTVKPGNFNKDFYDFAVSKGFFKEGSSVELKHFFFAQDGNAAISDYDYYRYNLPTMCDIRQLLFVVEVVDKYFDCYDGGCINANQVI'
        end_time = time()
        max_memory_used = 0
        return cleaned_seq, end_time - start_time, max_memory_used, 0
    elif property in enzymes:
        input_prompt = '[Generate ' + property.lower() + ' protein] ' + 'Seq=<' + sequence
    else:
        input_prompt = '[Generate ' + property + ' protein] ' + 'Seq=<' + sequence



    start_time = time()
    protein_seq = generator(input_prompt, temperature=0.5,
                            top_k=40,
                            top_p=0.9,
                            do_sample=True,
                            repetition_penalty=1.2,
                            max_new_tokens=seq_length, 
                            num_return_sequences=1)[0]["generated_text"]

    end_time = time()

    start_idx = protein_seq.find('Seq=<')
    end_idx = protein_seq.find('>', start_idx)
    protein_seq = protein_seq[start_idx:end_idx]
    cleaned_seq = clean_protein_sequence(protein_seq)
    tokens = tokenizer.encode(cleaned_seq, add_special_tokens=False)
    tokens_per_second = len(tokens) / (end_time - start_time)

    return cleaned_seq, end_time - start_time, tokens_per_second



# Create the Gradio interface

with gr.Blocks() as demo:
    gr.Markdown('''
    
    ### Interactive protein sequence generation and visualization. 

    Generating novel protein sequences possessing desired properties, termed as protein engineering, is crucial for industries like drug development and chemical synthesis.

    This model supports two types of generation, uncontrollable and controllable. Uncontrollable generation refers to generating any viable protein sequence, whereas controllable refers to generating proteins having a desired property or a characteristic. 

    ### Usage

    For uncontrollable generation, input any starting amino acids and press 'Submit' without choosing a property. For controllable generation, choose any of the ten properties supported by this model before pressing 'Submit'. The important inference metrics will be displayed along with the generated output. 

    ### Example

    As an example, the protein sequence corresponding to the SARS-CoV-2 Spike Protein is given. This example does not run inference on the model as the sequence is hard-coded, however it showcases the flow of interacting with the demo. 
    ''')
    
    with gr.Row():
        input_text = gr.Textbox(label="Enter starting amino acids for protein sequence generation", placeholder="Example input: MK")
    
    with gr.Row():
        seq_length = gr.Slider(2, 200, value=30, step=1, label="Length", info="Choose the number of tokens to generate")
        classes = ["SARS-CoV-2 Spike Protein (example)", 'Tetratricopeptide-like helical domain superfamily', 'CheY-like superfamily', 'S-adenosyl-L-methionine-dependent methyltransferase superfamily', 'Thioredoxin-like superfamily', "Non-Hemolytic" ,"Soluble", "Oxidoreductase", "Transferase", "Hydrolase", "Lyase", "Isomerase", "Ligase", "Translocase"]
        protein_property = gr.Dropdown(classes, label="Class")
    
    with gr.Row():
        btn = gr.Button("Submit")
    
    with gr.Row():
        output_text = gr.Textbox(label="Generated protein sequence will appear here")

    with gr.Row():

        infer_time = gr.Number(label="Inference Time (s)", precision=2)
        tokens_per_sec = gr.Number(label="Tokens/sec", precision=2)
    
    # with gr.Row():
    #     btn_vis = gr.Button("Visualize")
    
    # with gr.Row():
    #     structure_visual = gr.HTML()
    
    btn.click(generate_protein_sequence, inputs=[input_text, seq_length, protein_property], outputs=[output_text, infer_time, tokens_per_sec])

    # btn_vis.click(predict_structure, inputs=output_text, outputs=[structure_visual])


# Run the Gradio interface
demo.launch()