| import random |
| import torch |
| import gradio as gr |
| from gradio_rangeslider import RangeSlider |
| import pandas as pd |
| from utils import create_vocab, setup_seed |
| from dataset_mlm import get_paded_token_idx_gen, add_tokens_to_vocab |
| import time |
|
|
| is_stopped = False |
|
|
| def temperature_sampling(logits, temperature): |
| logits = logits / temperature |
| probabilities = torch.softmax(logits, dim=-1) |
| sampled_token = torch.multinomial(probabilities, 1) |
| return sampled_token |
|
|
| def stop_generation(): |
| global is_stopped |
| is_stopped = True |
| return "Generation stopped." |
|
|
| def CTXGen(τ, g_num, length_range, model_name, seed): |
| if seed =='random': |
| seed = random.randint(0,100000) |
| setup_seed(seed) |
| else: |
| setup_seed(int(seed)) |
| |
| global is_stopped |
| is_stopped = False |
|
|
| device = torch.device("cpu") |
| vocab_mlm = create_vocab() |
| vocab_mlm = add_tokens_to_vocab(vocab_mlm) |
| save_path = model_name |
| train_seqs = pd.read_csv('C0_seq.csv') |
| train_seq = train_seqs['Seq'].tolist() |
| model = torch.load(save_path, map_location=torch.device('cpu')) |
| model = model.to(device) |
|
|
| start, end = length_range |
| X1 = "X" |
| X2 = "X" |
| X4 = "" |
| X5 = "" |
| X6 = "" |
| model.eval() |
| with torch.no_grad(): |
| IDs = [] |
| generated_seqs = [] |
| generated_seqs_FINAL = [] |
| cls_pos_all = [] |
| cls_probability_all = [] |
| act_pos_all = [] |
| act_probability_all = [] |
|
|
| count = 0 |
| gen_num = int(g_num) |
| NON_AA = ["B", "O", "U", "Z", "X", '<K16>', '<α1β1γδ>', '<Ca22>', '<AChBP>', '<K13>', '<α1BAR>', '<α1β1ε>', '<α1AAR>', '<GluN3A>', '<α4β2>', |
| '<GluN2B>', '<α75HT3>', '<Na14>', '<α7>', '<GluN2C>', '<NET>', '<NavBh>', '<α6β3β4>', '<Na11>', '<Ca13>', |
| '<Ca12>', '<Na16>', '<α6α3β2>', '<GluN2A>', '<GluN2D>', '<K17>', '<α1β1δε>', '<GABA>', '<α9>', '<K12>', |
| '<Kshaker>', '<α3β4>', '<Na18>', '<α3β2>', '<α6α3β2β3>', '<α1β1δ>', '<α6α3β4β3>', '<α2β2>', '<α6β4>', '<α2β4>', |
| '<Na13>', '<Na12>', '<Na15>', '<α4β4>', '<α7α6β2>', '<α1β1γ>', '<NaTTXR>', '<K11>', '<Ca23>', |
| '<α9α10>', '<α6α3β4>', '<NaTTXS>', '<Na17>', '<high>', '<low>', '[UNK]', '[SEP]', '[PAD]', '[CLS]', '[MASK]'] |
|
|
| start_time = time.time() |
| while count < gen_num: |
| new_seq = None |
| if is_stopped: |
| return "output.csv", pd.DataFrame() |
|
|
| if time.time() - start_time > 1200: |
| break |
|
|
| gen_len = random.randint(int(start), int(end)) |
| X3 = "X" * gen_len |
| seq = [f"{X1}|{X2}|{X3}|{X4}|{X5}|{X6}"] |
| vocab_mlm.token_to_idx["X"] = 4 |
|
|
| padded_seq, _, _, _ = get_paded_token_idx_gen(vocab_mlm, seq, new_seq) |
| input_text = ["[MASK]" if i == "X" else i for i in padded_seq] |
|
|
| gen_length = len(input_text) |
| length = gen_length - sum(1 for x in input_text if x != '[MASK]') |
|
|
| for i in range(length): |
| if is_stopped: |
| return "output.csv", pd.DataFrame() |
|
|
| _, idx_seq, idx_msa, attn_idx = get_paded_token_idx_gen(vocab_mlm, seq, new_seq) |
| idx_seq = torch.tensor(idx_seq).unsqueeze(0).to(device) |
| idx_msa = torch.tensor(idx_msa).unsqueeze(0).to(device) |
| attn_idx = torch.tensor(attn_idx).to(device) |
|
|
| mask_positions = [j for j in range(gen_length) if input_text[j] == "[MASK]"] |
| mask_position = torch.tensor([mask_positions[torch.randint(len(mask_positions), (1,))]]) |
|
|
| logits = model(idx_seq, idx_msa, attn_idx) |
| mask_logits = logits[0, mask_position.item(), :] |
|
|
| predicted_token_id = temperature_sampling(mask_logits, τ) |
| predicted_token = vocab_mlm.to_tokens(int(predicted_token_id)) |
| input_text[mask_position.item()] = predicted_token |
| padded_seq[mask_position.item()] = predicted_token.strip() |
| new_seq = padded_seq |
|
|
| generated_seq = input_text |
|
|
| generated_seq[1] = "[MASK]" |
| generated_seq[2] = "[MASK]" |
| input_ids = vocab_mlm.__getitem__(generated_seq) |
| logits = model(torch.tensor([input_ids]).to(device), idx_msa) |
|
|
| cls_mask_logits = logits[0, 1, :] |
| act_mask_logits = logits[0, 2, :] |
|
|
| cls_probability, cls_mask_probs = torch.topk((torch.softmax(cls_mask_logits, dim=-1)), k=1) |
| act_probability, act_mask_probs = torch.topk((torch.softmax(act_mask_logits, dim=-1)), k=1) |
|
|
| cls_pos = vocab_mlm.idx_to_token[cls_mask_probs[0].item()] |
| act_pos = vocab_mlm.idx_to_token[act_mask_probs[0].item()] |
|
|
| cls_probability = cls_probability[0].item() |
| act_probability = act_probability[0].item() |
| generated_seq = generated_seq[generated_seq.index('[MASK]') + 2:generated_seq.index('[SEP]')] |
|
|
| if generated_seq.count('C') % 2 == 0 and len("".join(generated_seq)) == gen_len: |
| generated_seqs.append("".join(generated_seq)) |
| if "".join(generated_seq) not in train_seq and "".join(generated_seq) not in generated_seqs[0:-1] and all(x not in NON_AA for x in generated_seq): |
| generated_seqs_FINAL.append("".join(generated_seq)) |
| cls_pos_all.append(cls_pos) |
| cls_probability_all.append(cls_probability) |
| act_pos_all.append(act_pos) |
| act_probability_all.append(act_probability) |
| IDs.append(count+1) |
| out = pd.DataFrame({ |
| 'ID':IDs, |
| 'Generated_seq': generated_seqs_FINAL, |
| 'Subtype': cls_pos_all, |
| 'Subtype_probability': cls_probability_all, |
| 'Potency': act_pos_all, |
| 'Potency_probability': act_probability_all, |
| 'random_seed': seed |
| }) |
| out.to_csv("output.csv", index=False, encoding='utf-8-sig') |
| count += 1 |
| yield "output.csv", out |
| return "output.csv", out |
|
|
| with gr.Blocks() as demo: |
| gr.Markdown("🔗 **[Label Prediction](https://huggingface.co/spaces/oucgc1996/CreoPep_Label_Prediction)**" |
| "🔗 **[Unconstrained Generation](https://huggingface.co/spaces/oucgc1996/CreoPep_Unconstrained_generation)**" |
| "🔗 **[Conditional Generation](https://huggingface.co/spaces/oucgc1996/CreoPep_conditional_generation)**" |
| "🔗 **[Optimization Generation](https://huggingface.co/spaces/oucgc1996/CreoPep_optimization_generation)**") |
| gr.Markdown("# Conotoxin Unconstrained Generation") |
| gr.Markdown("#### Input") |
| gr.Markdown("✅**τ**: temperature factor controls the diversity of conotoxins generated. The higher the value, the higher the diversity.") |
| gr.Markdown("✅**Number of generations**: if it is not completed within 1200 seconds, it will automatically stop.") |
| gr.Markdown("✅**Length range**: expected length range of conotoxins generated.") |
| gr.Markdown("✅**Model**: model parameters trained at different stages of data augmentation. Please refer to the paper for details.") |
| gr.Markdown("✅**Seed**: enter an integer as the random seed to ensure reproducible results. The default is random.") |
| with gr.Row(): |
| τ = gr.Slider(minimum=1, maximum=2, step=0.1, label="τ") |
| g_num = gr.Dropdown(choices=[1, 10, 20, 30, 40, 50], label="Number of generations") |
| length_range = RangeSlider(minimum=8, maximum=50, step=1, value=(12, 16), label="Length range") |
| model_name = gr.Dropdown(choices=['model_final.pt','model_C1.pt','model_C2.pt','model_C3.pt','model_C4.pt','model_C5.pt','model_mlm.pt'], label="Model") |
| seed = gr.Textbox(label="Seed", value="random") |
| with gr.Row(): |
| start_button = gr.Button("Start Generation") |
| stop_button = gr.Button("Stop Generation") |
| with gr.Row(): |
| output_file = gr.File(label="Download generated conotoxins") |
| with gr.Row(): |
| output_df = gr.DataFrame(label="Generated Conotoxins") |
|
|
| start_button.click(CTXGen, inputs=[τ, g_num, length_range, model_name, seed], outputs=[output_file, output_df]) |
| stop_button.click(stop_generation, outputs=None) |
|
|
| demo.launch() |
