Update app.py

app.py

@@ -37,14 +37,16 @@ def process_sequence(sequence, domain_bounds, n):
                 logits = model(**inputs).logits
             mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]
             mask_token_logits = logits[0, mask_token_index, :]
+
+            # filter out non-amino acid tokens
+            filtered_indices = list(range(4, 23 + 1))
+            filtered_logits = mask_token_logits[:, filtered_indices]
             # Decode top n tokens
-            top_n_tokens = torch.topk(mask_token_logits, n, dim=1).indices[0].tolist()
+            top_n_tokens = torch.topk(filtered_logits, n, dim=1).indices[0].tolist()
             mutation = [tokenizer.decode([token]) for token in top_n_tokens]
             top_n_mutations[(sequence[i], i)] = mutation 
     
             logits_array = mask_token_logits.cpu().numpy()
-            # filter out non-amino acid tokens
-            filtered_indices = list(range(4, 23 + 1))
             filtered_logits = logits_array[:, filtered_indices]
             all_logits.append(filtered_logits)
     
@@ -66,7 +68,7 @@ def process_sequence(sequence, domain_bounds, n):
     plt.ylabel('Token')
     plt.xlabel('Residue Index')
     plt.yticks(rotation=0)
-    plt.xticks(x_tick_positions - start_index, x_tick_labels, rotation=0) 
+    plt.xticks(x_tick_positions - start_index + 0.5, x_tick_labels, rotation=0) 
     
     # Save the figure to a BytesIO object
     buf = BytesIO()