Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import pandas as pd
 import torch
 from transformers import AutoTokenizer, AutoModelForMaskedLM
+import torch.nn.functional as F
 import logging
 import numpy as np 
 import matplotlib.pyplot as plt
@@ -52,15 +53,15 @@ def process_sequence(sequence, domain_bounds, n):
     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())
+    normalized_logits_array = F.softmax(torch.tensor(all_logits_array), dim=-1).numpy()
     transposed_logits_array = normalized_logits_array.T
 
     # Plotting the heatmap
-    y_tick_positions = np.arange((start_index+1), end_index + 1, 10)
-    y_tick_labels = [str(pos) for pos in y_tick_positions]
+    x_tick_positions = np.arange(start_index, end_index)
+    x_tick_labels = [str(pos + 1) for pos in x_tick_positions]
     
     plt.figure(figsize=(15, 8))
-    sns.heatmap(transposed_logits_array, cmap='plasma', xticklabels=y_tick_labels, yticklabels=filtered_tokens)
+    sns.heatmap(transposed_logits_array, cmap='plasma', xticklabels=x_tick_labels, yticklabels=filtered_tokens)
     plt.title('Logits for masked per residue tokens')
     plt.ylabel('Token')
     plt.xlabel('Residue Index')