Update app.py

app.py

@@ -28,9 +28,62 @@ def generate_caption(protein, prompt):
     #     f.write('>{}\n'.format("protein_name"))
     #     f.write('{}\n'.format(protein.strip()))
     # os.system("python esm_scripts/extract.py esm2_t36_3B_UR50D /home/user/app/example.fasta /home/user/app --repr_layers 36 --truncation_seq_length 1024 --include per_tok")
-    esm_emb = run_demo(protein_name='protein_name', protein_seq=protein, 
-                       model=model_esm, alphabet=alphabet, 
-                       include='per_tok', repr_layers=[36], truncation_seq_length=1024)
+    # esm_emb = run_demo(protein_name='protein_name', protein_seq=protein, 
+    #                    model=model_esm, alphabet=alphabet, 
+    #                    include='per_tok', repr_layers=[36], truncation_seq_length=1024)
+    protein_name='protein_name'
+    protein_seq=protein
+    include='per_tok'
+    repr_layers=[36]
+    truncation_seq_length=1024
+    toks_per_batch=4096
+    
+    dataset = FastaBatchedDataset([protein_name], [protein_seq])
+    batches = dataset.get_batch_indices(toks_per_batch, extra_toks_per_seq=1)
+    data_loader = torch.utils.data.DataLoader(
+        dataset, collate_fn=alphabet.get_batch_converter(truncation_seq_length), batch_sampler=batches
+    )
+    print(f"Read sequences")
+    return_contacts = "contacts" in include
+    assert all(-(model_esm.num_layers + 1) <= i <= model_esm.num_layers for i in repr_layers)
+    repr_layers = [(i + model_esm.num_layers + 1) % (model_esm.num_layers + 1) for i in repr_layers]
+    with torch.no_grad():
+        for batch_idx, (labels, strs, toks) in enumerate(data_loader):
+            print(
+                f"Processing {batch_idx + 1} of {len(batches)} batches ({toks.size(0)} sequences)"
+            )
+            if torch.cuda.is_available():
+                toks = toks.to(device="cuda", non_blocking=True)
+            out = model_esm(toks, repr_layers=repr_layers, return_contacts=return_contacts)
+            logits = out["logits"].to(device="cpu")
+            representations = {
+                layer: t.to(device="cpu") for layer, t in out["representations"].items()
+            }
+            if return_contacts:
+                contacts = out["contacts"].to(device="cpu")
+            for i, label in enumerate(labels):
+                result = {"label": label}
+                truncate_len = min(truncation_seq_length, len(strs[i]))
+                # Call clone on tensors to ensure tensors are not views into a larger representation
+                # See https://github.com/pytorch/pytorch/issues/1995
+                if "per_tok" in include:
+                    result["representations"] = {
+                        layer: t[i, 1 : truncate_len + 1].clone()
+                        for layer, t in representations.items()
+                    }
+                if "mean" in include:
+                    result["mean_representations"] = {
+                        layer: t[i, 1 : truncate_len + 1].mean(0).clone()
+                        for layer, t in representations.items()
+                    }
+                if "bos" in include:
+                    result["bos_representations"] = {
+                        layer: t[i, 0].clone() for layer, t in representations.items()
+                    }
+                if return_contacts:
+                    result["contacts"] = contacts[i, : truncate_len, : truncate_len].clone()
+            esm_emb = result['representations'][36]
+    
     print("esm embedding generated")
     esm_emb = F.pad(esm_emb.t(), (0, 1024 - len(esm_emb))).t().to('cuda')
     print("esm embedding processed")