generation.py

# Copyright (c) 2023, Tri Dao, Albert Gu.

import argparse
import sys
import os
from pathlib import Path
import torch
import torch.nn.functional as F
current_directory = Path(__file__).parent.absolute()
parent_dir = current_directory.parent
print(parent_dir)
sys.path.append(str(current_directory))

import utils
from einops import rearrange
from utils.evo.tokenization import Vocab


def generate_aa_codon_map(amino_acid_to_codon,rna_vocab):
    aa_to_codon_indices = {}
    for aa, codons in amino_acid_to_codon.items():
        aa_to_codon_indices[aa] = []
        for codon in codons:
            aa_to_codon_indices[aa].append(rna_vocab.encode(codon)[1:-1].tolist())
    return aa_to_codon_indices

def CodonSplitter(seq):
    kmer_list = []
    for j in range(0, (len(seq) - 3) + 1,3):
        kmer_list.append(seq[j: j + 3])
    return kmer_list

def generate_sequence(model,rna_vocab,aa_to_codon,protein_id,args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    device = "cuda"
    model.to(device)
    model.eval()
    vocab_dict = rna_vocab.to_dict()
   
    SPECIAL_TOKEN_IDS = [
    rna_vocab.tokens_to_idx['<cls>'], 
    rna_vocab.tokens_to_idx['<unk>'],
    rna_vocab.tokens_to_idx['<pad>'],
    rna_vocab.tokens_to_idx['<mask>'],
    rna_vocab.tokens_to_idx["---"],
    rna_vocab.tokens_to_idx["NNN"],
    rna_vocab.tokens_to_idx["<eos>"],
    ]

    STOP_TOKEN_IDS = {
        rna_vocab.tokens_to_idx['UAA'],
        rna_vocab.tokens_to_idx['UAG'],
        rna_vocab.tokens_to_idx['UGA']
    }
    print(f"Number of parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")


    torch.random.manual_seed(0)
    if args.prompt is None:
        input_ids = torch.randint(0, 71, (args.batch, args.promptlen), dtype=torch.long, device="cuda")
        #print(input_ids.shape)
        protein_condition = args.protein_condition
        if protein_condition is not None:
            protein_condition = protein_condition.unsqueeze(0)
            #print(protein_condition.shape)
            protein_condition = protein_condition.to(device)
    else:
        input_ids = rna_vocab.encode(CodonSplitter(args.prompt))
        input_ids = torch.tensor(input_ids, dtype=torch.long, device="cuda")
        input_ids = torch.tensor(input_ids[:1], dtype=torch.long, device="cuda")#use the cls token as prompt
        input_ids = input_ids.unsqueeze(0)
        protein_condition = args.protein_condition
        if protein_condition is not None:
            protein_condition = protein_condition.to(device)
            protein_condition = protein_condition.unsqueeze(0)
    max_length = input_ids.shape[1] + args.genlen

    special_token_ids = [rna_vocab.tokens_to_idx[t] for t in ['<cls>', '<unk>', '<pad>', '<mask>', '---', 'NNN', '<eos>']]
    # 人类核基因组密码子表（不含 STOP）
    AA_TO_CODONS_NO_STOP = {aa: codons for aa, codons in aa_to_codon.items() if aa != '*'}

    stop_codons = {rna_vocab.tokens_to_idx['UAA'], rna_vocab.tokens_to_idx['UAG'], rna_vocab.tokens_to_idx['UGA']}

    fn = lambda: model.generate(
        input_ids=input_ids,
        aa_seq=args.aa_seq,
        embedding=protein_condition,#args.protein_condition,
        max_length=max_length,
        cg=True,
        return_dict_in_generate=True,
        output_scores=True,
        enable_timing=False,
        temperature=args.temperature,
        top_k=args.topk,
        top_p=args.topp,
        min_p=args.minp,
        repetition_penalty=args.repetition_penalty,
        special_token_ids=SPECIAL_TOKEN_IDS,
        aa_to_codon=aa_to_codon,
        tokenizer_vocab = vocab_dict,
        use_mask_predict=True,
        mask_predict_steps=args.mask_predict_steps,#default 12
        fill_schedule="linear",
        sampling_method=args.sampling_method,
        bias_alpha=args.bias_alpha,
        bias_e_power=args.bias_e_power,
        beam_size=args.beam_size,
        beam_candidates_per_pos=args.beam_candidates_per_pos,
        codon_usage_table_path = args.codon_usage_table_path,
        raw_beta = args.raw_beta,
        raw_top_n = args.raw_top_n,
        raw_min_w = args.raw_min_w,
        log_path= args.log_path,
        run_tag= args.run_tag,
    )

    out = fn()
    seq_ids = out
    out_tokens = seq_ids[0]  # 不要再 [1:]
    out_sequences = rna_vocab.decode(out_tokens.cpu().numpy())
    if args.bias_e_power is not None:
        outpath = Path(f"{args.outdir}/{args.model_name}_{args.sampling_method}_topk{args.topk:g}_topp{args.topp:g}_temp{args.temperature:g}_bias_e_power{args.bias_e_power:g}.fasta")
    else:
        outpath = Path(f"{args.outdir}/{args.model_name}_{args.sampling_method}_topk{args.topk:g}_topp{args.topp:g}_temp{args.temperature:g}_mask_predict_steps{args.mask_predict_steps:g}_beta{args.raw_beta}.fasta")
    outpath.parent.mkdir(parents=True, exist_ok=True)
    with open(f"{outpath}", "a") as f:
        f.write(f">{protein_id}\n")
        f.write(out_sequences +"\n")

    #return out_sequences

if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")