Update app.py

app.py

@@ -21,13 +21,13 @@ modelfile = 'model.pt'
 layers = 6
 heads = 16
 embed_dim = 128
-batch_toks = 1024
+batch_toks = 4096
 
 inp_len = 50
 
 device = "cpu"
 
-alphabet = Alphabet(standard_toks = 'AGCT')
+alphabet = Alphabet(mask_prob = 0.0, standard_toks = 'AGCT')
 assert alphabet.tok_to_idx == {'<pad>': 0, '<eos>': 1, '<unk>': 2, 'A': 3, 'G': 4, 'C': 5, 'T': 6, '<cls>': 7, '<mask>': 8, '<sep>': 9}
 
 class CNN_linear(nn.Module):
@@ -97,7 +97,7 @@ def eval_step(dataloader, model, threshold=0.5):
             # toks = toks.to(device)
             my_bar.progress((i+1)/len(dataloader), text="Running UTR_LM")
             # print(toks)
-            logits = model(toks)
+            logits = model(toks, return_representation = True, return_contacts=True)
 
             logits = logits.reshape(-1)
             # y_prob = torch.sigmoid(logits)
@@ -108,8 +108,8 @@ def eval_step(dataloader, model, threshold=0.5):
             # y_pred_list.extend(y_pred.tolist())
     
     st.success('Done', icon="✅")
-    # data_pred = pd.DataFrame({'ID':ids_list, 'Sequence':strs_list, "Translation Efficiency":logits_list, "prob":y_prob_list, "pred":y_pred_list})
-    data_pred = pd.DataFrame({'ID':ids_list, 'Sequence':strs_list, "Translation Efficiency":logits_list})
+    # data_pred = pd.DataFrame({'ID':ids_list, 'Sequence':strs_list, "MRL":logits_list, "prob":y_prob_list, "pred":y_pred_list})
+    data_pred = pd.DataFrame({'ID':ids_list, 'Sequence':strs_list, "MRL":logits_list})
     return data_pred
 
 def read_raw(raw_input):
@@ -135,7 +135,7 @@ def generate_dataset_dataloader(ids, seqs):
     dataset = FastaBatchedDataset(ids, seqs, mask_prob = 0.0)
     
     # dataset = FastaBatchedDataset(ids, seqs)
-    batches = dataset.get_batch_indices(toks_per_batch=batch_toks, extra_toks_per_seq=2)
+    batches = dataset.get_batch_indices(toks_per_batch=batch_toks, extra_toks_per_seq=1)
     dataloader = torch.utils.data.DataLoader(dataset, 
                                             collate_fn=alphabet.get_batch_converter(), 
                                             batch_sampler=batches, 
@@ -150,8 +150,10 @@ def predict_raw(raw_input):
     _, dataloader = generate_dataset_dataloader(ids, seqs)
 
     model = CNN_linear()
-
-    model.load_state_dict(torch.load(modelfile, map_location=torch.device('cpu')), strict = False)
+    print(model.state_dict().keys())
+    print(torch.load(modelfile, map_location=torch.device('cpu')).keys())
+    model.esm2.load_state_dict({k.replace('module.', ''):v for k,v in torch.load(modelfile, map_location=torch.device('cpu')).items()}, strict = False)
+    # model.load_state_dict(torch.load(modelfile, map_location=torch.device('cpu')), strict = False)
 
     # st.write('====Predict====')
     pred = eval_step(dataloader, model)