Reduce need for scaffold functions

app.py

@@ -20,12 +20,12 @@ theme = gr.themes.Monochrome(
 def get_model(model_name, token):
     login(token=token)
 
-    # if torch.cuda.is_available():
-    #     model = ESM3.from_pretrained(model_name, device=torch.device("cuda"))
-    # else:
-    #     model = ESM3.from_pretrained(model_name, device=torch.device("cpu"))
+    if torch.cuda.is_available():
+        model = ESM3.from_pretrained(model_name, device=torch.device("cuda"))
+    else:
+        model = ESM3.from_pretrained(model_name, device=torch.device("cpu"))
 
-    model = ESM3.from_pretrained(model_name, device=torch.device("cpu"))
+    # model = ESM3.from_pretrained(model_name, device=torch.device("cpu"))
     return model
 
 ## Function to render 3D structure using py3Dmol
@@ -48,46 +48,14 @@ def get_pdb(pdb_id, chain_id):
     return pdb
 
 
-# def select_motif(pdb, motif_start, motif_end):
-#     motif_inds = np.arange(motif_start, motif_end)
-#     motif_sequence = pdb[motif_inds].sequence
-#     motif_atom37_positions = pdb[motif_inds].atom37_positions
-#     return [motif_sequence, motif_atom37_positions]
-
-# def setup_prompt(prompt_length, motif_sequence, motif_atom37_positions, insert_size):
-#     prompt_length = 200
-
-#     sequence_prompt = ["_"]*prompt_length
-#     sequence_prompt[insert_size:insert_size+len(motif_sequence)] = list(motif_sequence)
-#     sequence_prompt = "".join(sequence_prompt)
-
-#     structure_prompt = torch.full((prompt_length, 37, 3), np.nan)
-#     structure_prompt[insert_size:insert_size+len(motif_atom37_positions)] = torch.tensor(motif_atom37_positions)
-
-#     protein_prompt = ESMProtein(sequence=sequence_prompt, coordinates=structure_prompt)
-
-#     return [sequence_prompt, structure_prompt, protein_prompt]
-
-
-# def generate_scaffold_sequence(model_name, token, sequence_prompt, protein_prompt):
-#     sequence_generation_config = GenerationConfig(track="sequence",
-#                                                   num_steps=sequence_prompt.count("_") // 2,
-#                                                   temperature=0.5)
-#     model = get_model(model_name, token)
-#     sequence_generation = model.generate(protein_prompt, sequence_generation_config)
-#     return sequence_generation
-
-
 def scaffold(model_name, token, pdb_id, chain_id, motif_start, motif_end, prompt_length, insert_size):
     pdb = get_pdb(pdb_id, chain_id)
-    # motif_sequence, motif_atom37_positions = select_motif(pdb, motif_start, motif_end)
 
+    ## Get motif sequence and atom37 positions
     motif_inds = np.arange(motif_start, motif_end)
     motif_sequence = pdb[motif_inds].sequence
     motif_atom37_positions = pdb[motif_inds].atom37_positions
 
-    # sequence_prompt, structure_prompt, protein_prompt = setup_prompt(prompt_length, motif_sequence, motif_atom37_positions, insert_size)
-
     ## Create sequence prompt
     sequence_prompt = ["_"]*prompt_length
     sequence_prompt[insert_size:insert_size+len(motif_sequence)] = list(motif_sequence)
@@ -97,10 +65,8 @@ def scaffold(model_name, token, pdb_id, chain_id, motif_start, motif_end, prompt
     structure_prompt = torch.full((prompt_length, 37, 3), np.nan)
     structure_prompt[insert_size:insert_size+len(motif_atom37_positions)] = torch.tensor(motif_atom37_positions)
 
-    ## Create protein prompt
+    ## Create protein prompt and sequence generation config
     protein_prompt = ESMProtein(sequence=sequence_prompt, coordinates=structure_prompt)
-
-    # sequence_generation = generate_scaffold_sequence(model_name, token, sequence_prompt, protein_prompt)
     sequence_generation_config = GenerationConfig(track="sequence",
                                                   num_steps=sequence_prompt.count("_") // 2,
                                                   temperature=0.5)
@@ -261,8 +227,12 @@ with gr.Blocks(theme=theme) as esm_app:
         gr.Markdown(
             """
             # ESM3: A frontier language model for biology.
-            - Created By: [EvolutionaryScale](https://www.evolutionaryscale.ai/blog/esm3-release)
-            - Spaces App By: [Tuple, The Cloud Genomics Company](https://tuple.xyz) [[Colby T. Ford](https://colbyford.com)]
+            Model Created By: [EvolutionaryScale](https://www.evolutionaryscale.ai)
+            - Press Release: https://www.evolutionaryscale.ai/blog/esm3-release
+            - GitHub: https://github.com/evolutionaryscale/esm
+            - HuggingFace Model: https://huggingface.co/EvolutionaryScale/esm3-sm-open-v1
+
+            Spaces App By: [Tuple, The Cloud Genomics Company](https://tuple.xyz) [[Colby T. Ford](https://colbyford.com)]
             """
         )
     with gr.Row():