Update

README.md

@@ -1,12 +1,18 @@
 ---
+# https://huggingface.co/docs/hub/spaces-config-reference
 title: SVM
-emoji: 🔥
-colorFrom: purple
-colorTo: yellow
-sdk: streamlit
-sdk_version: 1.21.0
+emoji: 🧬
+colorFrom: green
+colorTo: green
+sdk: gradio
 app_file: app.py
 pinned: false
+models:
+ - InstaDeepAI/nucleotide-transformer-500m-1000g
+ - facebook/esmfold_v1
+ - sentence-transformers/all-mpnet-base-v2
+python_version: 3.10.4
+license: mit
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,76 +1,275 @@
+# credit: https://huggingface.co/spaces/simonduerr/3dmol.js/blob/main/app.py
+from typing import Tuple
+import os
+import sys
+from urllib import request
+
+import gradio as gr
+import requests
+from transformers import AutoTokenizer, AutoModelForMaskedLM, EsmModel, AutoModel
 import torch
-import streamlit as st
-from transformers import AutoTokenizer, OPTForCausalLM
-
-
-@st.cache_resource
-def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("facebook/galactica-30b")
-    model = OPTForCausalLM.from_pretrained("facebook/galactica-30b", device_map='auto', low_cpu_mem_usage=True, torch_dtype=torch.float16)
-    model.gradient_checkpointing_enable()
-    return tokenizer, model
-
-
-st.set_page_config(
-    page_title='BioML-SVM',
-    layout="wide"
-)
-
-with st.spinner("Loading Models and Tokens..."):
-    tokenizer, model = load_model()
-
-with st.form(key='my_form'):
-    col1, col2 = st.columns([10, 1])
-    text_input = col1.text_input(label='Enter the amino sequence')
-    with col2:
-        st.text('')
-        st.text('')
-        submit_button = st.form_submit_button(label='Submit')
-
-    if submit_button:
-        st.session_state['result_done'] = False
-    # input_text = "[START_AMINO]GHMQSITAGQKVISKHKNGRFYQCEVVRLTTETFYEVNFDDGSFSDNLYPEDIVSQDCLQFGPPAEGEVVQVRWTDGQVYGAKFVASHPIQMYQVEFEDGSQLVVKRDDVYTLDEELP[END_AMINO]"
-        with st.spinner('Generating...'):
-            # formatted_text = f"[START_AMINO]{text_input}[END_AMINO]"
-            # formatted_text = f"Here is the sequence: [START_AMINO]{text_input}[END_AMINO]"
-            formatted_text = f"{text_input}"
-            input_ids = tokenizer(formatted_text, return_tensors="pt").input_ids.to("cuda")
-            outputs = model.generate(
-                input_ids=input_ids,
-                max_new_tokens=500
-            )
-            result = tokenizer.decode(outputs[0]).replace(formatted_text, "")
-        st.markdown(result)
-
-        if 'result_done' not in st.session_state or not st.session_state.result_done:
-            st.session_state['result_done'] = True
-            st.session_state['previous_state'] = result
+import progres as pg
+import esm
+
+import msa
+
+
+tokenizer_nt = AutoTokenizer.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-1000g")
+model_nt = AutoModelForMaskedLM.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-1000g")
+model_nt.eval()
+
+tokenizer_aa = AutoTokenizer.from_pretrained("facebook/esm2_t12_35M_UR50D")
+model_aa = EsmModel.from_pretrained("facebook/esm2_t12_35M_UR50D")
+model_aa.eval()
+
+tokenizer_se = AutoTokenizer.from_pretrained('sentence-transformers/all-mpnet-base-v2')
+model_se = AutoModel.from_pretrained('sentence-transformers/all-mpnet-base-v2')
+model_se.eval()
+
+msa_transformer, msa_transformer_alphabet = esm.pretrained.esm_msa1b_t12_100M_UR50S()
+msa_transformer = msa_transformer.eval()
+msa_transformer_batch_converter = msa_transformer_alphabet.get_batch_converter()
+
+
+
+def nt_embed(sequence: str):
+    tokens_ids = tokenizer_nt.batch_encode_plus([sequence], return_tensors="pt")["input_ids"]
+    attention_mask = tokens_ids != tokenizer_nt.pad_token_id
+    with torch.no_grad():
+        torch_outs = model_nt(
+            tokens_ids,#.to('cuda'),
+            attention_mask=attention_mask,#.to('cuda'),
+            output_hidden_states=True
+        )
+    last_layer_CLS = torch_outs.hidden_states[-1].detach()[:, 0, :][0]
+    return last_layer_CLS
+
+
+def aa_embed(sequence: str):
+    tokens = tokenizer_aa([sequence], return_tensors="pt")
+    with torch.no_grad():
+        torch_outs = model_aa(**tokens)
+    return torch_outs[0]
+
+
+def se_embed(sentence: str):
+    encoded_input = tokenizer_se([sentence], return_tensors='pt')
+    with torch.no_grad():
+        model_output = model_se(**encoded_input)
+    return model_output[0]
+
+
+def msa_embed(sequences: list):
+    inputs = msa.greedy_select(sequences, num_seqs=128) # can change this to pass more/fewer sequences
+    msa_transformer_batch_labels, msa_transformer_batch_strs, msa_transformer_batch_tokens = msa_transformer_batch_converter([inputs])
+    msa_transformer_batch_tokens = msa_transformer_batch_tokens.to(next(msa_transformer.parameters()).device)
+    
+    with torch.no_grad():
+        temp = msa_transformer(msa_transformer_batch_tokens,repr_layers=[12])['representations']
+    temp = temp[12][:,:,0,:]
+    temp = torch.mean(temp,(0,1))
+    return temp
+
+
+def go_embed(terms):
+    pass
+
+
+def download_data_if_required():
+    url_base = f"https://zenodo.org/record/{pg.zenodo_record}/files"
+    fps = [pg.trained_model_fp]
+    urls = [f"{url_base}/trained_model.pt"]
+    #for targetdb in pre_embedded_dbs:
+    #    fps.append(os.path.join(database_dir, targetdb + ".pt"))
+    #    urls.append(f"{url_base}/{targetdb}.pt")
+
+    if not os.path.isdir(pg.trained_model_dir):
+        os.makedirs(pg.trained_model_dir)
+    #if not os.path.isdir(database_dir):
+    #    os.makedirs(database_dir)
+
+    printed = False
+    for fp, url in zip(fps, urls):
+        if not os.path.isfile(fp):
+            if not printed:
+                print("Downloading data as first time setup (~340 MB) to ", pg.progres_dir,
+                      ", internet connection required, this can take a few minutes",
+                      sep="", file=sys.stderr)
+                printed = True
+            try:
+                request.urlretrieve(url, fp)
+                d = torch.load(fp, map_location="cpu")
+                if fp == pg.trained_model_fp:
+                    assert "model" in d
+                else:
+                    assert "embeddings" in d
+            except:
+                if os.path.isfile(fp):
+                    os.remove(fp)
+                print("Failed to download from", url, "and save to", fp, file=sys.stderr)
+                print("Exiting", file=sys.stderr)
+                sys.exit(1)
+
+    if printed:
+        print("Data downloaded successfully", file=sys.stderr)
+
+
+def get_pdb(pdb_code="", filepath=""):
+    if pdb_code is None or pdb_code == "":
+        try:
+            with open(filepath.name) as f:
+                return f.read()
+        except AttributeError as e:
+            return None
     else:
-        if 'result_done' in st.session_state and st.session_state.result_done:
-            st.markdown(st.session_state.previous_state)
-
-if 'result_done' in st.session_state and st.session_state.result_done:
-    with st.form(key='ask_more'):
-        col1, col2 = st.columns([10, 1])
-        text_input = col1.text_input(label='Ask more question')
-        with col2:
-            st.text('')
-            st.text('')
-            submit_button = st.form_submit_button(label='Submit')
-
-        if submit_button:
-            with st.spinner('Generating...'):
-                # formatted_text = f"[START_AMINO]{text_input}[END_AMINO]"
-                formatted_text = f"Q:{text_input}\n\nA:\n\n"
-                input_ids = tokenizer(formatted_text, return_tensors="pt").input_ids.to("cuda")
-
-                outputs = model.generate(
-                    input_ids=input_ids,
-                    max_length=len(formatted_text) + 500,
-                    do_sample=True,
-                    top_k=40,
-                    num_beams=1,
-                    num_return_sequences=1
+        return requests.get(f"https://files.rcsb.org/view/{pdb_code}.pdb").content.decode()
+
+
+def molecule(pdb):
+
+    x = (
+        """<!DOCTYPE html>
+        <html>
+        <head>    
+    <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
+    <style>
+    body{
+        font-family:sans-serif
+    }
+    .mol-container {
+    width: 100%;
+    height: 600px;
+    position: relative;
+    }
+    .mol-container select{
+        background-image:None;
+    }
+    </style>
+     <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js" integrity="sha512-STof4xm1wgkfm7heWqFJVn58Hm3EtS31XFaagaa8VMReCXAkQnJZ+jEy8PCC/iT18dFy95WcExNHFTqLyp72eQ==" crossorigin="anonymous" referrerpolicy="no-referrer"></script>
+    <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
+    </head>
+    <body>  
+    <div id="container" class="mol-container"></div>
+  
+            <script>
+               let pdb = `"""
+        + pdb
+        + """`  
+      
+             $(document).ready(function () {
+                let element = $("#container");
+                let config = { backgroundColor: "black" };
+                let viewer = $3Dmol.createViewer(element, config);
+                viewer.addModel(pdb, "pdb");
+                viewer.getModel(0).setStyle({}, { cartoon: { color:"spectrum" } });
+                viewer.addSurface("MS", { opacity: .5, color: "white" });
+                viewer.zoomTo();
+                viewer.render();
+                viewer.zoom(0.8, 2000);
+              })
+        </script>
+        </body></html>"""
+    )
+
+    return f"""<iframe style="width: 100%; height: 600px" name="result" allow="midi; geolocation; microphone; camera; 
+    display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
+    allow-scripts allow-same-origin allow-popups 
+    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
+    allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
+
+
+def str2coords(s):
+    coords = []
+    for line in s.split('\n'):
+        if (line.startswith("ATOM  ") or line.startswith("HETATM")) and line[12:16].strip() == "CA":
+            coords.append([float(line[30:38]), float(line[38:46]), float(line[46:54])])
+        elif line.startswith("ENDMDL"):
+            break
+    return coords
+
+
+def update_st(inp, file):
+    pdb = get_pdb(inp, file)
+    return (molecule(pdb), pg.embed_coords(str2coords(pdb)))
+
+
+def update_nt(inp):
+    return str(nt_embed(inp or ''))
+
+
+def update_aa(inp):
+    return str(aa_embed(inp))
+
+
+def update_se(inp):
+    return str(se_embed(inp))
+
+
+def update_go(inp):
+    return str(go_embed(inp))
+
+
+def update_msa(inp):
+    return str(msa_embed(msa.read_msa(inp.name)))
+
+
+demo = gr.Blocks()
+
+with demo:
+    with gr.Tabs():
+        with gr.TabItem("PDB Structural Embeddings"):
+            with gr.Row():
+                with gr.Box():
+                    inp = gr.Textbox(
+                        placeholder="PDB Code or upload file below", label="Input structure"
+                    )
+                    file = gr.File(file_count="single")
+                    gr.Examples(["2CBA", "6VXX"], inp)
+                    btn = gr.Button("View structure")
+            gr.Markdown("# PDB viewer using 3Dmol.js")
+            mol = gr.HTML()
+            emb = gr.Textbox(interactive=False)
+            btn.click(fn=update_st, inputs=[inp, file], outputs=[mol, emb])
+        with gr.TabItem("Nucleotide Sequence Embeddings"):
+            with gr.Box():
+                inp = gr.Textbox(
+                    placeholder="ATCGCTGCCCGTAGATAATAAGAGACACTGAGGCC", label="Input Nucleotide Sequence"
+                )
+                btn = gr.Button("View embeddings")
+                emb = gr.Textbox(interactive=False)
+                btn.click(fn=update_nt, inputs=[inp], outputs=emb)
+        with gr.TabItem("Amino Acid Sequence Embeddings"):
+            with gr.Box():
+                inp = gr.Textbox(
+                    placeholder="AAGQCYRGRCSGGLCCSKYGYCGSGPAYCG", label="Input Amino Acid Sequence"
+                )
+                btn = gr.Button("View embeddings")
+                emb = gr.Textbox(interactive=False)
+                btn.click(fn=update_aa, inputs=[inp], outputs=emb)
+        with gr.TabItem("Sentence Embeddings"):
+            with gr.Box():
+                inp = gr.Textbox(
+                    placeholder="Your text here", label="Input Sentence"
                 )
-                result = tokenizer.decode(outputs[0]).replace(formatted_text, "")
-            st.markdown(result)
+                btn = gr.Button("View embeddings")
+                emb = gr.Textbox(interactive=False)
+                btn.click(fn=update_se, inputs=[inp], outputs=emb)
+        with gr.TabItem("MSA Embeddings"):
+            with gr.Box():
+                inp = gr.File(file_count="single", label="Input MSA")
+                btn = gr.Button("View embeddings")
+                emb = gr.Textbox(interactive=False)
+                btn.click(fn=update_msa, inputs=[inp], outputs=emb)
+        with gr.TabItem("GO Embeddings"):
+            with gr.Box():
+                inp = gr.Textbox(
+                    placeholder="", label="Input GO Terms"
+                )
+                btn = gr.Button("View embeddings")
+                emb = gr.Textbox(interactive=False)
+                btn.click(fn=update_go, inputs=[inp], outputs=emb)
+
+
+if __name__ == "__main__":
+    download_data_if_required()
+    demo.launch()

gitattributes.txt

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

msa.py

@@ -0,0 +1,62 @@
+import glob
+import itertools
+from pathlib import Path
+from typing import List, Tuple, Optional, Dict, NamedTuple, Union, Callable
+import string
+
+import numpy as np
+import torch
+from scipy.spatial.distance import squareform, pdist, cdist
+from Bio import SeqIO
+#import biotite.structure as bs
+#from biotite.structure.io.pdbx import PDBxFile, get_structure
+#from biotite.database import rcsb
+from tqdm import tqdm
+import pandas as pd
+
+
+# This is an efficient way to delete lowercase characters and insertion characters from a string
+deletekeys = dict.fromkeys(string.ascii_lowercase)
+deletekeys["."] = None
+deletekeys["*"] = None
+translation = str.maketrans(deletekeys)
+
+
+def read_sequence(filename: str) -> Tuple[str, str]:
+    """ Reads the first (reference) sequences from a fasta or MSA file."""
+    record = next(SeqIO.parse(filename, "fasta"))
+    return record.description, str(record.seq)
+
+def remove_insertions(sequence: str) -> str:
+    """ Removes any insertions into the sequence. Needed to load aligned sequences in an MSA. """
+    return sequence.translate(translation)
+
+def read_msa(filename: str) -> List[Tuple[str, str]]:
+    """ Reads the sequences from an MSA file, automatically removes insertions."""
+    return [(record.description, remove_insertions(str(record.seq))) for record in SeqIO.parse(filename, "fasta")]
+
+
+def greedy_select(msa: List[Tuple[str, str]], num_seqs: int, mode: str = "max") -> List[Tuple[str, str]]:
+    """
+    Select sequences from the MSA to maximize the hamming distance
+    Alternatively, can use hhfilter 
+    """
+    assert mode in ("max", "min")
+    if len(msa) <= num_seqs:
+        return msa
+    
+    array = np.array([list(seq) for _, seq in msa], dtype=np.bytes_).view(np.uint8)
+
+    optfunc = np.argmax if mode == "max" else np.argmin
+    all_indices = np.arange(len(msa))
+    indices = [0]
+    pairwise_distances = np.zeros((0, len(msa)))
+    for _ in range(num_seqs - 1):
+        dist = cdist(array[indices[-1:]], array, "hamming")
+        pairwise_distances = np.concatenate([pairwise_distances, dist])
+        shifted_distance = np.delete(pairwise_distances, indices, axis=1).mean(0)
+        shifted_index = optfunc(shifted_distance)
+        index = np.delete(all_indices, indices)[shifted_index]
+        indices.append(index)
+    indices = sorted(indices)
+    return [msa[idx] for idx in indices]

requirements.txt

@@ -1,5 +1,13 @@
-transformers
 accelerate
-streamlit
-# bitsandbytes
-# scipy
+gradio==3.33.1
+--find-links https://data.pyg.org/whl/torch-2.0.0+cpu.html pyg-lib==0.2.0+pt20
+requests==2.31.0
+torch==2.0.1
+--find-links https://data.pyg.org/whl/torch-2.0.0+cpu.html torch-cluster==1.6.1
+torch-geometric==2.3.1
+torch-scatter==2.1.1
+--find-links https://data.pyg.org/whl/torch-2.0.0+cpu.html torch-sparse==0.6.17
+--find-links https://data.pyg.org/whl/torch-2.0.0+cpu.html torch-spline-conv==1.2.2
+transformers==4.29.2
+progres
+fair-esm