generate_mamba_condition.py

import argparse
import time
import json
import random
from collections import Counter
from dataset import Mol3DDataset, SimpleTokenizer
import argparse
from utils import set_seed
import numpy as np
import wandb
import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.nn import functional as F
from mamba import MambaLMHeadModel, MambaConfig
from functools import partial
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel

def load_tokenizer(tokenizer_path,max_length):
    tokenizer = SimpleTokenizer(max_length)  # Update max_length if needed
    tokenizer.load_vocab(tokenizer_path)
    return tokenizer

def get_first_token_distribution(train_data_path):
    with open(train_data_path, 'r') as file:
        data = file.readlines()
    first_tokens = [line.strip().split()[0] for line in data if len(line.strip())]
    token_counts = Counter(first_tokens)
    total = sum(token_counts.values())
    token_probs = {token: count / total for token, count in token_counts.items()}
    return token_probs


def sample_first_token(token_probs):
    tokens, probs = zip(*token_probs.items())
    first_token = random.choices(tokens, weights=probs)[0]
    return first_token


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--run_name', type=str,
                        help="name for wandb run", required=False)
    parser.add_argument('--debug', action='store_true',
                        default=False, help='debug')
    # in moses dataset, on average, there are only 5 molecules per scaffold
    parser.add_argument('--scaffold', action='store_true',
                        default=False, help='condition on scaffold')
    parser.add_argument('--lstm', action='store_true',
                        default=False, help='use lstm for transforming scaffold')
    parser.add_argument('--data_name', type=str, default='moses2',
                        help="name of the dataset to train on", required=False)
    # parser.add_argument('--property', type=str, default = 'qed', help="which property to use for condition", required=False)
    parser.add_argument('--props', nargs="+", default=['qed'],
                        help="properties to be used for condition", required=False)
    parser.add_argument('--num_props', type=int, default = 0, help="number of properties to use for condition", required=False)
    # parser.add_argument('--prop1_unique', type=int, default = 0, help="unique values in that property", required=False)
    parser.add_argument('--n_layer', type=int, default=8,
                        help="number of layers", required=False)
    parser.add_argument('--n_head', type=int, default=8,
                        help="number of heads", required=False)
    parser.add_argument('--n_embd', type=int, default=768,
                        help="embedding dimension", required=False)
    parser.add_argument('--max_epochs', type=int, default=30,
                        help="total epochs", required=False)
    parser.add_argument('--batch_size', type=int, default=80,
                        help="batch size", required=False)
    parser.add_argument('--num_workers', type=int, default=8,
                        help="number of workers for data loaders", required=False)
    parser.add_argument('--learning_rate', type=float,
                        default=1e-4, help="learning rate", required=False)
    parser.add_argument('--lstm_layers', type=int, default=0,
                        help="number of layers in lstm", required=False)
    parser.add_argument('--max_len', type=int, default=512,
                        help="max_len", required=False)
    parser.add_argument('--epoch', type=int, default=0,
                        help="saved model epoch", required=False)
    parser.add_argument('--root_path', default='QM9_seq/spherical_seq',
                        help="Path to the root data directory", required=False)
    parser.add_argument('--output_tokenizer_dir', default='spherical_seq/tokenizer',
                        help="Path to the saved tokenizer directory", required=False)
    parser.add_argument('--conditions_path', default=None,
                        help="Path to the generation condition", required=False)
    parser.add_argument('--save_path', default='mamba_test',
                        help="Path to the generation file dir", required=False)
    parser.add_argument("--model-name", type=str, default="state-spaces/mamba-130m")
    parser.add_argument("--prompt", type=str, default=None)
    parser.add_argument("--promptlen", type=int, default=100)
    parser.add_argument("--genlen", type=int, default=100)
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--topk", type=int, default=50)
    parser.add_argument("--topp", type=float, default=1.0)
    parser.add_argument("--repetition-penalty", type=float, default=1.0)
    parser.add_argument("--batch", type=int, default=1)
    parser.add_argument("--repeats", type=int, default=12000)
    parser.add_argument("--seed", type=int, default=44)

    args = parser.parse_args()

    set_seed(int(args.seed))

    os.environ["WANDB_MODE"] = "dryrun"

    max_len = args.max_len
    tokenizer_path = args.output_tokenizer_dir + "/vocab.json"
    print(tokenizer_path)
    print("tokenizer:")
    tokenizer = load_tokenizer(tokenizer_path,max_len)
    print(tokenizer.get_vocab())  # Print vocabulary
    vocab_size = tokenizer.get_vocab_size()

    repeats = args.repeats
    device = "cuda"
    dtype = torch.float16

    print("loading model")
    # model = GPT(mconf)
    if args.epoch == 0:
        model_path = f'cond_gpt/weights/{args.run_name}.pt'
        gen_name = f'{args.run_name}'
    else:
        model_path = f'cond_gpt/weights/{args.run_name}_ep{args.epoch}.pt'
        gen_name = f'{args.run_name}_ep{args.epoch}_top{args.topk}_temp{args.temperature}'
    train_data_path = args.root_path + '.txt'
    if args.conditions_path is not None:
        conditions_path = args.conditions_path + '.txt'

    # Load model and tokenizer
    mamba_config = MambaConfig(d_model=args.n_embd, n_layer=args.n_layer, vocab_size=vocab_size,
                               num_props=args.num_props, scaffold=args.scaffold)
    model = MambaLMHeadModel(mamba_config)
    model.load_state_dict(torch.load(model_path, map_location='cpu'), strict=True)
    model.to(device)
    model.eval()
    print("loaded model")
    print(f"Number of parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")


    # Get first token distribution
    if args.conditions_path is not None:
        print("loading conditions")
        with open(args.conditions_path, 'r') as file:
            data = file.readlines()
        first_tokens = [line.strip() for line in data]
        print('loaded conditions')
    else:
        print("missing conditions")

    generate = partial(model.generate,
                 max_length=max_len,
                 cg=True,
                 return_dict_in_generate=True,
                 output_scores=True,
                 enable_timing=False,
                 temperature=args.temperature,
                 top_k=args.topk,
                 top_p=args.topp,
                 repetition_penalty=args.repetition_penalty,
                 )

    num_samples = args.repeats  # Number of samples to generate for each first token
    os.makedirs(args.save_path, exist_ok=True)
    print(f'Saving samples to: {args.save_path}')
    # Generate samples for each first token and write to separate files
    for idx, first_token in enumerate(first_tokens):
        first_token_ids = torch.tensor([tokenizer.generation_encode(first_token) for _ in range(num_samples)],
                                       dtype=torch.long).to(device)
        out = generate(
            first_token_ids)  # out.sequences, out.scores
        samples = [tokenizer.decode(ids) for ids in out.sequences.cpu()]
        print(f'Completed {idx} with {len(samples)} samples')
        print(samples[0])
        with open(f'{args.save_path}/gen_mol_{idx}.txt', 'w') as f:
            for sample in samples:
                if sample.startswith(first_token):
                    gen_sample = sample[len(first_token):]
                    f.write(gen_sample + '\n')
                else:
                    print(f"{sample} does not start with {first_token}")