Update app.py

app.py

@@ -1,10 +1,31 @@
-#@title Lauch the Interface
-#@markdown Run to launch the interface.
-#@markdown 1. Enter the protein sequence.
-#@markdown 2. Specify the start and end index (inclusive) of the domain for which you would like to predict mutations (note that indexing starts at 1).
-#@markdown 3. Select the number of tokens you would like the model to predict for each position.
-#@markdown 4. Click 'Submit'.
-#@markdown 5. Click 'Download Outputs' to download the zip file.
+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
+import seaborn as sns
+from io import BytesIO
+from PIL import Image
+from contextlib import contextmanager
+import warnings
+import sys
+import os
+import zipfile
+
+logging.getLogger("transformers.modeling_utils").setLevel(logging.ERROR)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+# Load the tokenizer and model
+model_name = "ChatterjeeLab/FusOn-pLM"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForMaskedLM.from_pretrained(model_name, trust_remote_code=True)
+model.to(device)
+model.eval()
+
 
 def process_sequence(sequence, domain_bounds, n):
     start_index = int(domain_bounds['start'][0]) - 1