generate_n_shot.py

import argparse
import json
import os
import functools
import datetime
import random
from typing import Optional, Tuple
from tqdm import tqdm

from transformers import SuppressTokensLogitsProcessor
from transformers.generation import MinNewTokensLengthLogitsProcessor

import torch
from torch.utils.data import DataLoader, DistributedSampler

from enhancerdiff.constants import ENHANCER_ALPHABET, END_AL, MSA_PAD, START_AL, START_UL, END_UL, DNA, SEP, GAP, STOP
from enhancerdiff.datasets import EnhancerDataset
from enhancerdiff.collators import EnhancARCollator

from dayhoff.utils import load_msa_config_and_model, seed_everything, load_checkpoint
import numpy as np


# default to a single-GPU setup if not present
RANK = int(os.environ.get("RANK", 0))
LOCAL_RANK = int(os.environ.get("LOCAL_RANK", 0))
WORLD_SIZE = int(os.environ.get("WORLD_SIZE", 1))
DEVICE = torch.device(f"cuda:{LOCAL_RANK}")
print("device", DEVICE)


def get_val_dataloader(config, tokenizer, args):
    # load the dataset
    print("Preparing MSA dataset", flush=True)
    ds_train = EnhancerDataset(
        args.data_fpath,
        args.set,
        'max_hamming',
        config["n_sequences"],
        config["max_seq_len"],
        min_depth=None,
        gap_fraction=config["gap_fraction"],
        indel=config["indel"],
        max_id_to_query=args.max_id_to_query,
        no_query=False,
        min_msa_after_filter=16
    )
    print("Preparing MSA collator", flush=True)

    collator = EnhancARCollator(
        tokenizer=tokenizer,
        pad_to_multiple_of=config["pad_to_multiple_of"],
        flip_prob=False,
        return_anchor=True,
        sort_by_length=config["sort_by_length"],
    )

    print("Preparing MSA loader", flush=True)
    sampler = DistributedSampler(ds_train, num_replicas=WORLD_SIZE, rank=RANK, shuffle=False, drop_last=False)
    dl = DataLoader(
        dataset=ds_train, batch_size=1, sampler=sampler, num_workers=8,
        collate_fn=collator, pin_memory=True
    )

    return dl


def generate(args: argparse.Namespace) -> None:
    runs = np.arange(0, args.n_samples)

    for run in runs:
        # print(f"Starting job on rank {RANK} with local rank {LOCAL_RANK} and world size {WORLD_SIZE}")
        seed_everything(int(run))
        suffix = '.fa.aln'
        new_suffix = "_" + str(run) + '.fa.aln'
        #print(f"Starting job on rank {RANK} with local rank {LOCAL_RANK} and world size {WORLD_SIZE}")
        seed_everything(args.random_seed)
        # dist.init_process_group(backend="nccl")
        #if args.verbose:
            #print("Initializing model...", RANK)
        torch.cuda.set_device(LOCAL_RANK)

        # load model parameters from config file
        config, tokenizer, model, block = load_msa_config_and_model(os.path.join(args.in_fpath, "config.json"),
                                                                    alphabet=ENHANCER_ALPHABET, use_flash_attention_2=True)
        #if args.verbose:
            #print("Done initializing model.", RANK)

        # Load model and optimizer onto CPU
        initial_epoch, total_steps, total_tokens, total_seqs, _ = load_checkpoint(
            model, None, None, args.in_fpath, args.checkpoint_step
        )
        model = model.eval()
        model = model.to(DEVICE)
        model = model.to(torch.bfloat16)

        if args.verbose:
            print(ENHANCER_ALPHABET)
            print(tokenizer.a_to_i)
            print(tokenizer.i_to_a)

        if args.set != "unconditional":
            max_len = config["max_seq_len"]
            loader = get_val_dataloader(config, tokenizer, args)

            if args.min_p is not None:
                out_dir = os.path.join(args.in_fpath, "generations_n_times", "ckpt%d_T%.1f_m%.1f_minp_%.2f_s%d"
                                       % (total_steps, args.temp, args.max_id_to_query, args.min_p, args.random_seed) + args.set)
            else:
                out_dir = os.path.join(args.in_fpath, "generations_n_times", "ckpt%d_T%.1f_m%.1f_s%d"
                                       % (total_steps, args.temp, args.max_id_to_query, args.random_seed) + args.set)
            if RANK == 0:
                os.makedirs(out_dir, exist_ok=True)

            all_tokens = list(range(16))
            allowed_tokens = [ENHANCER_ALPHABET.index(aa) for aa in DNA]
            if config['no_query']:
                eos_id = ENHANCER_ALPHABET.index(SEP)
                allowed_tokens += [ENHANCER_ALPHABET.index(SEP)]
            else:
                eos_id = ENHANCER_ALPHABET.index(STOP)
                allowed_tokens += [ENHANCER_ALPHABET.index(STOP)]

            model.module.generation_config.eos_token_id = eos_id
            sup = SuppressTokensLogitsProcessor([t for t in all_tokens if not t in allowed_tokens], device=DEVICE)

            for it, batch in enumerate(loader):
                filename = loader.dataset.filenames[(it * WORLD_SIZE + RANK) % len(loader.dataset)]
                anchor, src, tgt = batch
                src = src.to(DEVICE)

                # Get the end of the MSA and prune depending on if this is query or no-query
                _, gen_start_idx = ((src == tokenizer.a_to_i[END_UL]).nonzero(as_tuple=True))
                if len(gen_start_idx) == 1:
                    if config['no_query']:
                        sample = src[:, :gen_start_idx]
                    else:
                        sample = src[:, :gen_start_idx + 2]
                else:   # This is a failsafe if we can't identify the MSA end-token...
                    sample = src

                if sample.shape[1] == 0:
                    continue

                sample_untokenized = [tokenizer.untokenize(s) for s in sample][0]

                procs = [sup]
                if args.min_len:
                    # Get length of shortest sequence
                    s = tokenizer.untokenize(src[0])
                    msa = s.replace('[', '').replace(']', '').replace('@', '').replace('*', '').replace('{', '').replace('}', '')
                    msa = [x for x in msa.split('/') if x != '']
                    msa = msa[::-1]
                    min_len = min(map(len, msa)) * args.min_len
                    min_proc = MinNewTokensLengthLogitsProcessor(prompt_length_to_skip = len(sample), min_new_tokens=min_len,
                                                                 eos_token_id=eos_id, device=DEVICE)
                    procs.append(min_proc)

                generated = model.module.generate(sample, do_sample=True, logits_processor=procs,
                                                 temperature=args.temp, num_beams=1, max_new_tokens=max_len,
                                                  use_cache=True, min_p=args.min_p)
                #generated = generated[0, -max_len:]
                untokenized = tokenizer.untokenize(generated[0])
                untokenized = [x for x in untokenized.split('/') if x != ''][-1]    # get the last sequence, which should be generation
                anchor_untokenized = tokenizer.untokenize(anchor[0])

                if args.verbose:
                    print("Prompt %d " % (it + 1) + sample_untokenized, datetime.datetime.now(), flush=True)
                    print("Generation %d " % (it + 1) + untokenized, datetime.datetime.now(), flush=True)
                    print ("Anchor %d " % (it + 1) + anchor_untokenized, datetime.datetime.now(), flush=True)

                with open(os.path.join(out_dir, filename.replace(suffix, new_suffix)), "w") as f:
                    f.write(">generation\n" + str(untokenized) + "\n")
                    s = tokenizer.untokenize(src[0])
                    msa = s.replace('[', '').replace(']', '').replace('@', '').replace('*', '').replace('{', '').replace('}', '')
                    msa = [x for x in msa.split('/') if x != '']
                    f.write(">%d\n" % 0)
                    f.write(anchor_untokenized + "\n")
                    for i, seq in enumerate(msa[::-1]):
                        f.write(">%d\n" %(i+1))
                        f.write(seq + "\n")
        else:
            out_dir = os.path.join(args.in_fpath, "generations", "ckpt%d_T%.1f_s%d_unconditional"
                                   % (total_steps, args.temp, args.random_seed))
            if RANK == 0:
                os.makedirs(out_dir, exist_ok=True)

            all_tokens = list(range(16))
            allowed_tokens = [ENHANCER_ALPHABET.index(aa) for aa in DNA]
            if "gap" in args.in_fpath:
                start_seq = ENHANCER_ALPHABET.index(START_AL)
                eos_id = ENHANCER_ALPHABET.index(END_AL)
                allowed_tokens += [ENHANCER_ALPHABET.index(GAP)]
            elif "indel" in args.in_fpath:
                start_seq = ENHANCER_ALPHABET.index(START_UL)
                eos_id = ENHANCER_ALPHABET.index(END_UL)
            else:
                raise ValueError("Model should have gap or indel in name.")
            start = torch.tensor([[start_seq]]).to(DEVICE)
            sep_id = ENHANCER_ALPHABET.index(SEP)
            sep_tensor = torch.tensor([[sep_id]]).to(DEVICE)
            start = torch.repeat_interleave(start, args.batch_size, dim=0)
            sep_id = torch.repeat_interleave(sep_tensor, args.batch_size, dim=0)
            model.module.generation_config.eos_token_id = eos_id
            sup = SuppressTokensLogitsProcessor([t for t in all_tokens if not t in allowed_tokens], device=DEVICE)
            n_seqs = 32
            for s in tqdm(range(args.n_generations // args.batch_size)):
                generated = start.clone()
                for n in range(n_seqs):
                    generated = model.module.generate(generated, do_sample=True, logits_processor=[sup],
                                                      temperature=args.temp, num_beams=1, max_new_tokens=512,
                                                      use_cache=True, min_p=args.min_p)
                    generated = torch.cat([generated, sep_id.clone()], dim=1)
                untokenized = [tokenizer.untokenize(g) for g in generated]
                for n, unt in enumerate(untokenized):
                    n_gen = s * args.batch_size + n
                    with open(os.path.join(out_dir, 'rank%d_%d.fasta' %(RANK, n_gen)), "w") as f:
                        unt = unt[1:-1] # Strip out whatever stop and start
                        # Replace all things in the middle with new lines
                        for sep in [SEP]:
                            unt = unt.replace(sep, " ")
                        unt = unt.split()
                        for i, seq in enumerate(unt):
                            f.write(">%d\n" %i)
                            f.write(seq + "\n")
                            print(">%d" %i)
                            print(seq, flush=True)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("in_fpath", type=str)  # location of checkpoint
    parser.add_argument("data_fpath", type=str)  # location of test data
    parser.add_argument("--set", type=str, default="100_test_alignments")  # location of test data
    parser.add_argument("--max_id_to_query", type=float, default=1.0)  #
    parser.add_argument("--verbose", action="store_true")
    parser.add_argument("--checkpoint_step", type=int, default=-1)
    parser.add_argument("--temp", type=float, default=1.0)  #
    parser.add_argument("--random_seed", type=int, default=0)  #
    parser.add_argument("--batch_size", type=int, default=16)  #
    parser.add_argument("--n_generations", type=int, default=16*100000)  #
    parser.add_argument("--min_p", type=float, default=None)  #
    parser.add_argument("--min_len", type=float, default=None)
    parser.add_argument("--n_samples", type=int, default=10)  #


    args = parser.parse_args()
    generate(args)


if __name__ == "__main__":
    main()