fix app.py

app.py

@@ -1,8 +1,7 @@
 # app.py
 # --- IMPORTS ---
 import re
-from pathlib import Path
-
+import esm
 import gradio as gr
 import torch
 from torch.utils.data import DataLoader
@@ -28,11 +27,15 @@ def generate_smiles_for_sequence(protein_sequence: str, num_samples: int):
     if len(protein_sequence) < 10:
         raise gr.Error("Protein sequence is too short.")
 
-    embedding = get_esm_embedding(
-        protein_sequence,
-        'esm2_t33_650M_UR50D',
-        device
-    ).to(dtype=torch.bfloat16)
+    print(">> inference started, attempts:", num_samples, flush=True)
+
+    with torch.no_grad():
+        batch_converter = alphabet.get_batch_converter()
+        _, _, tokens = batch_converter([("protein", protein_sequence)])
+        tokens = tokens.to(device)
+        embedding = esm_model(tokens, repr_layers=[33])["representations"][33][:, 1:-1, :]
+        embedding = embedding.float() if device == "cpu" else embedding.bfloat16()
+
     n_batches = num_samples // 10
     dataset = InferenceDataset(embedding, batch_size=10, n_batches=n_batches)
     loader = DataLoader(dataset, batch_size=None)
@@ -56,6 +59,10 @@ def generate_smiles_for_sequence(protein_sequence: str, num_samples: int):
 
 # Load models on app startup
 device = "cuda" if torch.cuda.is_available() else "cpu"
+esm_model, alphabet = esm.pretrained.load_model_and_alphabet('esm2_t33_650M_UR50D')
+esm_model.eval()
+esm_model = esm_model.to(device)
+
 tokenizer_path = hf_hub_download(
     repo_id=REPO_ID,
     filename=TOKENIZER_FILENAME,
@@ -141,6 +148,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
 # Launch the app
 if __name__ == "__main__":
+    demo.queue(max_size=10)
     demo.launch(share=True)