import os import re import shutil import subprocess import sys import threading from itertools import groupby from pathlib import Path from tempfile import NamedTemporaryFile import gradio as gr import matplotlib.pyplot as plt import numpy as np import pandas as pd import spaces import torch from transformers import T5EncoderModel, T5Tokenizer # ----------------------------------------------------------------------------- # Reproducible upstream sources: these match the working TREAD Colab demo. # ----------------------------------------------------------------------------- TREAD_REPO = "https://github.com/KYQiu21/TREAD.git" TREAD_COMMIT = "7a7b3f571778cb89035c798c88df7eacb63ed2c1" PROTT5_MODEL = "Rostlab/prot_t5_xl_half_uniref50-enc" APP_DIR = Path(__file__).resolve().parent TREAD_DIR = APP_DIR / ".tread_source" # Web-demo guardrails only; these are not biological/model validity limits. MAX_GPU_SEQUENCE_LENGTH = int(os.getenv("TREAD_MAX_GPU_SEQUENCE_LENGTH", "1200")) MAX_CPU_SEQUENCE_LENGTH = int(os.getenv("TREAD_MAX_CPU_SEQUENCE_LENGTH", "600")) REPEAT_HEADS = [ "Any repeat (segmentation head)", "Alpha solenoid", "TIM-barrel", "Beta-propeller", "Beta-barrel", "Beta-solenoid", "Alpha/beta solenoid", ] TYPE_HEAD_INDEX = { "Alpha solenoid": 0, "TIM-barrel": 1, "Beta-propeller": 2, "Beta-barrel": 3, "Beta-solenoid": 4, "Alpha/beta solenoid": 5, } EXAMPLE_REPEAT = ( "MMKRNILAVIVPALLVAGTANAAEIYNKDGNKVDLYGKAVGLHYFSKGNGENSYGGNGDMTYARLGFKGETQINSDLTGYGQWEYNFQGNNSEGADAQTGNKTRLAFAGLKYADVGSFDYGRNYGVVYDALGYTDMLPEFGGDTAYSDDFFVGRVGGVATYRNSNFFGLVDGLNFAVQYLGKNERDTARRSNGDGVGGSISYEYEGFGIVGAYGAADRTNLQEAQPLGNGKKAEQWATGLKYDANNIYLAANYGETRNATPITNKFTNTSGFANKTQDVLLVAQYQFDFGLRPSIAYTKSKAKDVEGIGDVDLVNYFEVGATYYFNKNMSTYVDYIINQIDSDNKLGVGSDDTVAVGIVYQF" ) EXAMPLE_PROPELLER = ( "MEQVLYLGSYTKRESKGVHQIILDTDKKELRDYRLIAEVDSPTYLDLSADKGTLYSISKTDEGGGITSFKKNENGTYDKVAEISAEGSAPCYIHYDEDKKLIFTANYHGGYLTVYKENADGSFTMSDRAQHEGSSIHENQTIPHVHYSALSPDKKFLLACDLGTDEVYTYTVSDEGKLTEAARYKATPGTGPRHLVFHPNGKVAYLFGELSSDVEVLAYEAATGTFSLLQVITTIPAEHTGFNGGAAIRISADGKFVYASNRGHDSLVVYAVSEDGETLSLVEYVPTEGNTPRDFNLDPSGQFVIVAHQDSDNLTLFERDATTGKLTLVQKDVYAPECVCVFY" ) # ----------------------------------------------------------------------------- # Model state # # ZeroGPU path: # The GPU models are placed on "cuda" during app startup. In a ZeroGPU Space, # Hugging Face's CUDA emulation handles this even though a real GPU is only # assigned while a @spaces.GPU function is executing. # # CPU path: # A separate float32 copy is loaded lazily only if somebody presses the CPU # button. This avoids paying the RAM cost unless CPU fallback is actually used. # ----------------------------------------------------------------------------- _gpu_tokenizer = None _gpu_encoder = None _gpu_repeat_model = None _gpu_propeller_model = None _gpu_init_error = None _cpu_lock = threading.Lock() _cpu_tokenizer = None _cpu_encoder = None _cpu_repeat_model = None _cpu_propeller_model = None _cpu_init_error = None def _ensure_tread_source(): """Fetch exactly the same TREAD commit used by the public Colab demo.""" package_init = TREAD_DIR / "tread" / "__init__.py" if package_init.is_file(): if str(TREAD_DIR) not in sys.path: sys.path.insert(0, str(TREAD_DIR)) return shutil.rmtree(TREAD_DIR, ignore_errors=True) subprocess.run(["git", "init", str(TREAD_DIR)], check=True) subprocess.run( ["git", "-C", str(TREAD_DIR), "remote", "add", "origin", TREAD_REPO], check=True, ) subprocess.run( [ "git", "-C", str(TREAD_DIR), "fetch", "--depth", "1", "origin", TREAD_COMMIT, ], check=True, ) subprocess.run( ["git", "-C", str(TREAD_DIR), "checkout", "--detach", "FETCH_HEAD"], check=True, ) if not package_init.is_file(): raise FileNotFoundError( f"TREAD installation incomplete: {package_init} not found" ) if str(TREAD_DIR) not in sys.path: sys.path.insert(0, str(TREAD_DIR)) def _build_repeat_model(DMDModel, device): model = DMDModel( per_resi_emb_dim=1024, out_channel=64, hidden_dim=64, num_block=2, dropout=0.2, bilstm=True, kernel_size_conv1=3, kernel_size_block=7, multi=True, num_types=6, device=device, ) state = torch.load( TREAD_DIR / "trained_model" / "linear-edge_model_repeatsdb.pt", map_location="cpu", weights_only=True, ) model.load_state_dict(state) model = model.to(device).eval() return model def _build_propeller_model(DMDModel, device): model = DMDModel( per_resi_emb_dim=1024, out_channel=128, hidden_dim=64, num_block=1, dropout=0.2, bilstm=True, kernel_size_conv1=11, kernel_size_block=7, device=device, ) state = torch.load( TREAD_DIR / "trained_model" / "linear-edge_model_propeller_blade.pt", map_location="cpu", weights_only=True, ) model.load_state_dict(state) model = model.to(device).eval() return model def initialize_gpu_models_at_startup(): """Prepare the ZeroGPU model copy on emulated CUDA during app startup.""" global _gpu_tokenizer, _gpu_encoder, _gpu_repeat_model global _gpu_propeller_model, _gpu_init_error try: _ensure_tread_source() from tread.model import DMDModel _gpu_tokenizer = T5Tokenizer.from_pretrained( PROTT5_MODEL, do_lower_case=False, legacy=True, ) _gpu_encoder = T5EncoderModel.from_pretrained(PROTT5_MODEL) _gpu_encoder = _gpu_encoder.to("cuda").eval() _gpu_repeat_model = _build_repeat_model(DMDModel, torch.device("cuda")) _gpu_propeller_model = _build_propeller_model( DMDModel, torch.device("cuda") ) print("ZeroGPU model copy initialized successfully.") except Exception as exc: # Do not kill the whole web app: CPU fallback may still work and the # error will be shown clearly if the ZeroGPU button is pressed. _gpu_init_error = ( f"{type(exc).__name__}: {exc}" ) print("ZeroGPU initialization failed:", _gpu_init_error) def initialize_cpu_models(): """Load a separate CPU copy only when the CPU fallback is requested.""" global _cpu_tokenizer, _cpu_encoder, _cpu_repeat_model global _cpu_propeller_model, _cpu_init_error if all( x is not None for x in ( _cpu_tokenizer, _cpu_encoder, _cpu_repeat_model, _cpu_propeller_model, ) ): return if _cpu_init_error is not None: raise RuntimeError(_cpu_init_error) with _cpu_lock: if all( x is not None for x in ( _cpu_tokenizer, _cpu_encoder, _cpu_repeat_model, _cpu_propeller_model, ) ): return try: _ensure_tread_source() from tread.model import DMDModel # Tokenizers are device independent. Reuse the one already loaded # for ZeroGPU if available. _cpu_tokenizer = _gpu_tokenizer if _cpu_tokenizer is None: _cpu_tokenizer = T5Tokenizer.from_pretrained( PROTT5_MODEL, do_lower_case=False, legacy=True, ) # CPU inference uses float32 for broad PyTorch CPU compatibility. # low_cpu_mem_usage reduces peak RAM while loading this large model. _cpu_encoder = T5EncoderModel.from_pretrained( PROTT5_MODEL, torch_dtype=torch.float32, low_cpu_mem_usage=True, ) _cpu_encoder = _cpu_encoder.to("cpu").eval() _cpu_repeat_model = _build_repeat_model( DMDModel, torch.device("cpu") ) _cpu_propeller_model = _build_propeller_model( DMDModel, torch.device("cpu") ) print("CPU fallback model copy initialized successfully.") except Exception as exc: _cpu_init_error = f"{type(exc).__name__}: {exc}" raise RuntimeError(_cpu_init_error) from exc def clean_sequence(raw_sequence: str, max_length: int) -> str: if raw_sequence is None: raise gr.Error("Please paste a protein sequence.") text = raw_sequence.strip() if not text: raise gr.Error("Please paste a protein sequence.") # Accept either a raw sequence or a single FASTA record. lines = [line.strip() for line in text.splitlines() if line.strip()] if lines and lines[0].startswith(">"): lines = [line for line in lines[1:] if not line.startswith(">")] sequence = "".join(lines) sequence = re.sub(r"\s+", "", sequence).upper() if not sequence: raise gr.Error("No amino-acid sequence was found.") # Standard 20 amino acids plus the symbols handled by the notebook's # ProtT5 preprocessing (U, Z, O, B -> X) and X itself. allowed = set("ACDEFGHIKLMNPQRSTVWYUZOBX") invalid = sorted(set(sequence) - allowed) if invalid: raise gr.Error( "Invalid sequence characters: " + ", ".join(invalid) + ". Paste amino-acid letters only (a FASTA header is allowed)." ) if len(sequence) > max_length: raise gr.Error( f"This backend currently accepts up to {max_length} residues; " f"your sequence has {len(sequence)} residues." ) return sequence def get_prott5_embedding(sequence, tokenizer, encoder, device): processed = " ".join(list(re.sub(r"[UZOB]", "X", sequence))) ids = tokenizer( [processed], add_special_tokens=True, padding="longest", truncation=False, return_attention_mask=True, ) input_ids = torch.tensor(ids["input_ids"], device=device) attention_mask = torch.tensor(ids["attention_mask"], device=device) with torch.inference_mode(): embedding_repr = encoder( input_ids=input_ids, attention_mask=attention_mask, ) # Drop the terminal special token, exactly as in the working Colab demo. emb = embedding_repr.last_hidden_state[0, :-1] if emb.shape[0] != len(sequence): raise RuntimeError( f"ProtT5 returned {emb.shape[0]} residue embeddings for a " f"{len(sequence)}-residue sequence." ) return emb def get_ranges(preds, cutoff1=0.5, min_len=15, cutoff2=0.5, frac2=0.5): """Same motif-range logic as tread.utils.get_ranges.""" preds = np.asarray(preds).flatten() above_threshold = preds > cutoff1 peaks = [] for key, group in groupby(enumerate(above_threshold), key=lambda x: x[1]): if key: group = list(group) if len(group) >= min_len: beg = group[0][0] end = beg + len(group) if ( len(np.where(preds[beg:end] > cutoff2)[0]) / len(group) >= frac2 ): peaks.append((beg, end)) return peaks def make_plot(scores, ranges, title, cutoff): scores = np.asarray(scores).flatten() x = np.arange(1, len(scores) + 1) fig, ax = plt.subplots(figsize=(10, 4.5), dpi=150) ax.plot(x, scores, linewidth=1.5) ax.axhline( cutoff, linestyle="--", linewidth=1.0, alpha=0.7, label=f"cutoff = {cutoff:g}", ) for start0, end0 in ranges: ax.axvspan(start0 + 1, end0, alpha=0.18) ax.set_xlim(1, max(1, len(scores))) ax.set_ylim(0, 1) ax.set_xlabel("Residue") ax.set_ylabel("Residue score") ax.set_title(title) ax.legend(loc="upper right") fig.tight_layout() return fig def make_download_table(sequence, scores, ranges): scores = np.asarray(scores).flatten() predicted = np.zeros(len(sequence), dtype=bool) for start0, end0 in ranges: predicted[start0:end0] = True return pd.DataFrame( { "residue": np.arange(1, len(sequence) + 1), "amino_acid": list(sequence), "score": scores, "predicted_motif": predicted, } ) def save_csv(df: pd.DataFrame): temp = NamedTemporaryFile( prefix="tread_prediction_", suffix=".csv", delete=False, ) temp.close() df.to_csv(temp.name, index=False) return temp.name def _predict_core( sequence_text, model_choice, repeat_head, cutoff, min_len, *, tokenizer, encoder, repeat_model, propeller_model, device, backend_label, max_length, ): sequence = clean_sequence(sequence_text, max_length=max_length) embedding = get_prott5_embedding( sequence, tokenizer=tokenizer, encoder=encoder, device=device, ) with torch.inference_mode(): if model_choice == "RepeatsDB repeat annotation": seg_prediction, type_prediction = repeat_model.predict_single( embedding ) if repeat_head == "Any repeat (segmentation head)": scores = seg_prediction else: scores = type_prediction[TYPE_HEAD_INDEX[repeat_head]] profile_name = repeat_head title = f"TREAD RepeatsDB — {profile_name}" else: scores = propeller_model.predict_single(embedding) profile_name = "Beta-propeller blade" title = "TREAD — beta-propeller blade annotation" ranges = get_ranges( scores, cutoff1=float(cutoff), min_len=int(min_len), ) fig = make_plot(scores, ranges, title, float(cutoff)) # Convert notebook's 0-based, end-exclusive ranges to user-facing # 1-based inclusive residue coordinates. if ranges: ranges_df = pd.DataFrame( [ { "Start (1-based)": start0 + 1, "End (1-based)": end0, "Length": end0 - start0, } for start0, end0 in ranges ] ) range_text = ", ".join(f"{s + 1}–{e}" for s, e in ranges) else: ranges_df = pd.DataFrame( columns=["Start (1-based)", "End (1-based)", "Length"] ) range_text = "None at the current thresholds" profile_df = make_download_table(sequence, scores, ranges) csv_path = save_csv(profile_df) status = ( f"**Sequence length:** {len(sequence)} aa \n" f"**Model:** {model_choice} \n" f"**Profile:** {profile_name} \n" f"**Predicted motif ranges:** {range_text} \n" f"**Compute backend:** {backend_label}" ) return status, fig, ranges_df, csv_path @spaces.GPU(duration=120) def predict_gpu(sequence_text, model_choice, repeat_head, cutoff, min_len): """ZeroGPU path. A real GPU is allocated only for this function call.""" if _gpu_init_error is not None: raise gr.Error( "The ZeroGPU model copy could not be initialized. " f"Startup error: {_gpu_init_error}" ) if any( x is None for x in ( _gpu_tokenizer, _gpu_encoder, _gpu_repeat_model, _gpu_propeller_model, ) ): raise gr.Error( "ZeroGPU models are not initialized. Check the Space runtime log." ) device_name = "ZeroGPU" try: device_name = f"ZeroGPU — {torch.cuda.get_device_name(0)}" except Exception: pass return _predict_core( sequence_text, model_choice, repeat_head, cutoff, min_len, tokenizer=_gpu_tokenizer, encoder=_gpu_encoder, repeat_model=_gpu_repeat_model, propeller_model=_gpu_propeller_model, device=torch.device("cuda"), backend_label=device_name, max_length=MAX_GPU_SEQUENCE_LENGTH, ) def predict_cpu(sequence_text, model_choice, repeat_head, cutoff, min_len): """CPU fallback. Does not consume ZeroGPU quota, but is much slower.""" try: initialize_cpu_models() except Exception as exc: raise gr.Error(f"CPU model initialization failed: {exc}") from exc return _predict_core( sequence_text, model_choice, repeat_head, cutoff, min_len, tokenizer=_cpu_tokenizer, encoder=_cpu_encoder, repeat_model=_cpu_repeat_model, propeller_model=_cpu_propeller_model, device=torch.device("cpu"), backend_label="CPU fallback", max_length=MAX_CPU_SEQUENCE_LENGTH, ) def load_repeat_example(): return ( EXAMPLE_REPEAT, "RepeatsDB repeat annotation", "Beta-barrel", 0.8, 20, ) def load_propeller_example(): return ( EXAMPLE_PROPELLER, "Beta-propeller blade annotation", "Beta-propeller", 0.8, 20, ) # IMPORTANT for ZeroGPU: # Prepare the CUDA/emulated-CUDA model copy at module level, before requests. initialize_gpu_models_at_startup() with gr.Blocks(title="TREAD — Protein Repeat Annotation") as demo: gr.Markdown( """ # TREAD — Protein Repeat Annotation Paste a protein sequence and run one of the two pretrained TREAD models. - **RepeatsDB repeat annotation:** residue-wise repeat segmentation plus six repeat-fold heads. - **Beta-propeller blade annotation:** residue-wise blade prediction. **Recommended:** use **Run with ZeroGPU**. **Fallback:** use **Run on CPU** if GPU quota/availability is a problem; CPU inference is substantially slower. ProtT5 embeddings are generated on the fly with `Rostlab/prot_t5_xl_half_uniref50-enc`. """ ) with gr.Row(): with gr.Column(scale=3): sequence_input = gr.Textbox( label="Protein sequence", lines=10, placeholder=( "Paste a raw amino-acid sequence or a single FASTA record..." ), ) model_choice = gr.Radio( choices=[ "RepeatsDB repeat annotation", "Beta-propeller blade annotation", ], value="RepeatsDB repeat annotation", label="Model", ) repeat_head = gr.Dropdown( choices=REPEAT_HEADS, value="Beta-barrel", label="RepeatsDB profile to display", info=( "Used only for the RepeatsDB model. " "The default matches the public Colab example." ), ) with gr.Column(scale=2): cutoff = gr.Slider( minimum=0.0, maximum=1.0, value=0.8, step=0.01, label="Residue score threshold", ) min_len = gr.Slider( minimum=1, maximum=100, value=20, step=1, label="Minimum motif length", ) gpu_button = gr.Button( "Run with ZeroGPU (recommended)", variant="primary", ) cpu_button = gr.Button( "Run on CPU (slow fallback)", variant="secondary", ) with gr.Row(): repeat_example_button = gr.Button("Load RepeatsDB example") propeller_example_button = gr.Button("Load propeller example") gr.Markdown("## Results") status_output = gr.Markdown() plot_output = gr.Plot(label="Residue-wise score profile") ranges_output = gr.Dataframe( headers=["Start (1-based)", "End (1-based)", "Length"], label="Predicted motif ranges", interactive=False, ) download_output = gr.File(label="Download per-residue CSV") gr.Markdown( f""" ### Notes - ZeroGPU accepts sequences up to **{MAX_GPU_SEQUENCE_LENGTH} aa** in this web demo. - CPU fallback accepts sequences up to **{MAX_CPU_SEQUENCE_LENGTH} aa** by default because ProtT5-XL is very slow on the free CPU backend. - The range-calling logic uses the same thresholding rule as the public Colab/TREAD utility. - The web table reports **1-based inclusive** residue coordinates for readability. - Rare/ambiguous residues `U`, `Z`, `O`, and `B` are mapped to `X` for ProtT5 embedding, matching the Colab demo. Source code and pretrained TREAD checkpoints: [KYQiu21/TREAD](https://github.com/KYQiu21/TREAD) """ ) gpu_button.click( fn=predict_gpu, inputs=[ sequence_input, model_choice, repeat_head, cutoff, min_len, ], outputs=[ status_output, plot_output, ranges_output, download_output, ], ) cpu_button.click( fn=predict_cpu, inputs=[ sequence_input, model_choice, repeat_head, cutoff, min_len, ], outputs=[ status_output, plot_output, ranges_output, download_output, ], ) repeat_example_button.click( fn=load_repeat_example, inputs=[], outputs=[ sequence_input, model_choice, repeat_head, cutoff, min_len, ], ) propeller_example_button.click( fn=load_propeller_example, inputs=[], outputs=[ sequence_input, model_choice, repeat_head, cutoff, min_len, ], ) demo.queue() if __name__ == "__main__": demo.launch()