Update app.py

app.py

@@ -47,53 +47,53 @@ def process_sequence(sequence, domain_bounds, n):
             filtered_logits = logits_array[:, filtered_indices]
             all_logits.append(filtered_logits)
     
-        token_indices = torch.arange(logits.size(-1))
-        tokens = [tokenizer.decode([idx]) for idx in token_indices]
-        filtered_tokens = [tokens[i] for i in filtered_indices]
-        
-        all_logits_array = np.vstack(all_logits)
-        normalized_logits_array = (all_logits_array - all_logits_array.min()) / (all_logits_array.max() - all_logits_array.min())
-        transposed_logits_array = normalized_logits_array.T
+    token_indices = torch.arange(logits.size(-1))
+    tokens = [tokenizer.decode([idx]) for idx in token_indices]
+    filtered_tokens = [tokens[i] for i in filtered_indices]
+    
+    all_logits_array = np.vstack(all_logits)
+    normalized_logits_array = (all_logits_array - all_logits_array.min()) / (all_logits_array.max() - all_logits_array.min())
+    transposed_logits_array = normalized_logits_array.T
 
-        # Plotting the heatmap
-        step = 50
-        y_tick_positions = np.arange(0, len(sequence), step)
-        y_tick_labels = [str(pos) for pos in y_tick_positions]
-        
-        plt.figure(figsize=(15, 8))
-        sns.heatmap(transposed_logits_array, cmap='plasma', xticklabels=y_tick_labels, yticklabels=filtered_tokens)
-        plt.title('Logits for masked per residue tokens')
-        plt.ylabel('Token')
-        plt.xlabel('Residue Index')
-        plt.yticks(rotation=0)
-        plt.xticks(y_tick_positions, y_tick_labels, rotation = 0)
-        
-        # Save the figure to a BytesIO object
-        buf = BytesIO()
-        plt.savefig(buf, format='png')
-        buf.seek(0)
-        plt.close()
-        
-        # Convert BytesIO object to an image
-        img = Image.open(buf)
+    # Plotting the heatmap
+    step = 50
+    y_tick_positions = np.arange(0, len(sequence), step)
+    y_tick_labels = [str(pos) for pos in y_tick_positions]
     
-        original_residues = []
-        mutations = []
-        positions = []
+    plt.figure(figsize=(15, 8))
+    sns.heatmap(transposed_logits_array, cmap='plasma', xticklabels=y_tick_labels, yticklabels=filtered_tokens)
+    plt.title('Logits for masked per residue tokens')
+    plt.ylabel('Token')
+    plt.xlabel('Residue Index')
+    plt.yticks(rotation=0)
+    plt.xticks(y_tick_positions, y_tick_labels, rotation = 0)
     
-        for key, value in top_n_mutations.items():
-            original_residue, position = key
-            original_residues.append(original_residue)
-            mutations.append(value)
-            positions.append(position + 1)
-
-        df = pd.DataFrame({
-            'Original Residue': original_residues,
-            'Predicted Residues (in order of decreasing likelihood)': mutations,
-            'Position': positions
-        })
+    # Save the figure to a BytesIO object
+    buf = BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    plt.close()
     
-        df = df[start_index:end_index]
+    # Convert BytesIO object to an image
+    img = Image.open(buf)
+
+    original_residues = []
+    mutations = []
+    positions = []
+
+    for key, value in top_n_mutations.items():
+        original_residue, position = key
+        original_residues.append(original_residue)
+        mutations.append(value)
+        positions.append(position + 1)
+
+    df = pd.DataFrame({
+        'Original Residue': original_residues,
+        'Predicted Residues (in order of decreasing likelihood)': mutations,
+        'Position': positions
+    })
+
+    df = df[start_index:end_index]
             
     return df, img