import gradio as gr import pandas as pd import numpy as np import transformers as T import torch import matplotlib.pyplot as plt import shutil from torch.utils.data import DataLoader from utils import load_model,plot_tracks,calc_mutation_effect,plot_tracks_comparision from attention_extractor import attention_extractor import subprocess tokenizer = None model = None def load_tokenizer_and_model(): global tokenizer, model if tokenizer is None: tokenizer = T.BertTokenizer.from_pretrained("./model/") if model is None: load_tokenizer_and_model() model = load_model(tokenizer, "./model/model.bin") prefix_code = pd.read_csv('./data/prefix_codes.csv') prefix_code_dic = {k:v for k,v in zip(prefix_code.prefix,prefix_code.code_prefix)} rbp_options = prefix_code['RBP'].unique().tolist() def sequence_process(RBP, cell_line, ss): target = "_".join([RBP, cell_line]) barcode = prefix_code_dic[target] if len(ss) < 3: raise ValueError("Input sequence length is too short for 3-mer processing.") ss = [ss[i:i+3] for i in range(len(ss) - 2)] # 3 mer data seq = [barcode] seq.extend(ss[:-1]) return seq def predict_affinity(RBP, cell_line, ss): sequence = sequence_process(RBP, cell_line, ss) inputs = tokenizer(sequence, return_tensors="pt", padding=True, truncation=True, max_length=512) with torch.no_grad(): outputs = model(**inputs) affinity_score = outputs.logits.squeeze().item() plot_tracks(sequence, outputs.logits.squeeze()) plt.savefig("track_plot.png") plt.close() return affinity_score, "track_plot.png" def predict_affinity_with_mutation(RBP, cell_line, ss_wt, ss_mt): original_affinity = predict_affinity(RBP, cell_line, ss_wt) mutated_affinity = predict_affinity(RBP, cell_line, ss_mt) change_in_affinity = calc_mutation_effect(original_affinity, mutated_affinity) plot_tracks_comparision(ss_wt, ss_mt, original_affinity, mutated_affinity) plt.savefig("mutation_plot.png") plt.close() return change_in_affinity, "mutation_plot.png" def load_eclip_data(path): with open(path, 'r') as f: lines = f.readlines() try: assert np.all([i.startswith('>') for e, i in enumerate(lines) if e == 0 or e % 2 == 0]) except AssertionError: raise ValueError("Input file is not in the expected FASTA format.") seq = [i.strip() for i in lines if not i.startswith('>')] return np.array(seq) class SequenceDataset(torch.utils.data.Dataset): def __init__(self, barcode, seq, tokenizer, max_length=512): self.barcode = barcode self.tokenizer = tokenizer self.max_length = max_length self.sequence = seq self.n = len(self.sequence) def __len__(self): return self.n def __getitem__(self, i): ss = self.sequence[i] ss = [ss[i:i+3] for i in range(len(ss) - 2)] # 3 mer data seq = [self.barcode] seq.extend(ss[:-1]) inputs = self.tokenizer(seq, is_split_into_words=True, add_special_tokens=True, return_tensors='pt') return inputs['input_ids'] def extract_attention_regions(barcode, seq, layer, head): val_dataset = SequenceDataset(barcode, seq, tokenizer) val_loader = DataLoader( val_dataset, batch_size=4, shuffle = False, drop_last=False, num_workers=2 ) seq_attn = attention_extractor(model, val_loader, model.model.device, len(val_dataset)) if len(seq_attn) <= layer or len(seq_attn[layer]) <= head: raise ValueError("Invalid layer or head index for attention extraction.") attention_data = seq_attn[layer][head].squeeze() high_attention_regions = (attention_data > np.percentile(attention_data, 99)).nonzero(as_tuple=True) return high_attention_regions def predict_motif(RBP, cell_line, file, layer, head): target = "_".join([RBP, cell_line]) barcode = prefix_code_dic[target] seq = load_eclip_data(file) high_attention_regions = extract_attention_regions(barcode, seq, layer, head) with open("high_attention_regions.fasta", "w") as f: for i, (seq_idx, pos) in enumerate(zip(*high_attention_regions)): subseq = seq[seq_idx][pos:pos+10] # 10-mer f.write(f">{i}\n{subseq}\n") zip_file = "high_attention_regions.zip" shutil.make_archive("high_attention_regions", 'zip', ".", "high_attention_regions.fasta") return zip_file def single_sequence_affinity(RBP, cell_line, ss): return predict_affinity(RBP, cell_line, ss) def mutation_affinity(RBP, cell_line, ss_wt, ss_mt): return predict_affinity_with_mutation(RBP, cell_line, ss_wt, ss_mt) def motif_enrichment(RBP, cell_line, file, layer, head): return predict_motif(RBP, cell_line, file, layer, head) # defined input and output in Gradio def get_cell_lines(selected_rbp): filtered_df = prefix_code[prefix_code['RBP'] == selected_rbp] cell_lines = filtered_df['cell_line'].unique().tolist() return gr.update(choices=cell_lines) def build_gradio_interface(): rbp_dropdown = gr.Dropdown(label="Select RBP", choices=rbp_options, value=rbp_options[0]) cell_line_dropdown = gr.Dropdown(label="Select Cell Line", choices=[]) rbp_dropdown.change(fn=get_cell_lines, inputs=rbp_dropdown, outputs=cell_line_dropdown) affinity_interface = gr.Interface( fn=single_sequence_affinity, inputs=[ rbp_dropdown, cell_line_dropdown, gr.Textbox(label='Input RNA Sequence') ], outputs=[ gr.Textbox(label="Predicted Affinity"), gr.Image(label="Track Plot") ], description="Predicted binding affinity" ) mutation_interface = gr.Interface( fn=mutation_affinity, inputs=[ rbp_dropdown, cell_line_dropdown, gr.Textbox(label='Input Wild-type RNA Sequence'), gr.Textbox(label='Input Mutated RNA Sequence') ], outputs=[ gr.Textbox(label="Affinity Change"), gr.Image(label="Mutation Plot") ], description="Predicted mutation effect" ) motif_interface = gr.Interface( fn=motif_enrichment, inputs=[ rbp_dropdown, cell_line_dropdown, gr.File(label="Upload Sequence File"), gr.Slider(minimum=0, maximum=12, step=1, label="Select Attention Layer"), gr.Slider(minimum=0, maximum=11, step=1, label="Select Attention Head") ], outputs=gr.File(label="Download High Attention Regions (ZIP) (result.zip)"), description="Motif enrichment preparation: Download high attention regions and run AME manually." ) app = gr.TabbedInterface( interface_list=[affinity_interface, mutation_interface, motif_interface], tab_names=["Binding Affinity Prediction", "Mutation Effect Prediction", "Motif Enrichment"] ) return app if __name__ == "__main__": app = build_gradio_interface() app.launch(share=True)