Add files using upload-large-folder tool

08af434 verified 5 days ago

10.9 kB

from tqdm import tqdm

import json
import os

from typing import List, Dict, Tuple
import fsspec
import hydra
import lightning as L
import omegaconf
import rich.syntax
import rich.tree
import torch
import algo
import dataloader as dataloader
import utils
from text_metrics import evaluate_text_metrics


omegaconf.OmegaConf.register_new_resolver("cwd", os.getcwd)
omegaconf.OmegaConf.register_new_resolver("device_count", torch.cuda.device_count)
omegaconf.OmegaConf.register_new_resolver("eval", eval)
omegaconf.OmegaConf.register_new_resolver("div_up", lambda x, y: (x + y - 1) // y)

def _load_from_checkpoint(diffusion_model, config, tokenizer):
    if "hf" in config.algo.backbone:
        return diffusion_model(config, tokenizer=tokenizer).to("cuda")
    
    
    model = diffusion_model.load_from_checkpoint(
        config.eval.checkpoint_path, tokenizer=tokenizer, config=config
    )
    return model



@L.pytorch.utilities.rank_zero_only
def _print_config(
    config: omegaconf.DictConfig, resolve: bool = True, save_cfg: bool = True
) -> None:
    """Prints content of DictConfig using Rich library and its tree structure.
    Args:
      config (DictConfig): Configuration composed by Hydra.
      resolve (bool): Whether to resolve reference fields of DictConfig.
      save_cfg (bool): Whether to save the configuration tree to a file.
    """

    style = "dim"
    tree = rich.tree.Tree("CONFIG", style=style, guide_style=style)

    fields = config.keys()
    for field in fields:
        branch = tree.add(field, style=style, guide_style=style)

        config_section = config.get(field)
        branch_content = str(config_section)
        if isinstance(config_section, omegaconf.DictConfig):
            branch_content = omegaconf.OmegaConf.to_yaml(
                config_section, resolve=resolve
            )

        branch.add(rich.syntax.Syntax(branch_content, "yaml"))
    rich.print(tree)
    if save_cfg:
        with fsspec.open(
            "{}/config_tree.txt".format(config.checkpointing.save_dir), "w"
        ) as fp:
            rich.print(tree, file=fp)


@L.pytorch.utilities.rank_zero_only
def _print_batch(train_ds, valid_ds, tokenizer, k=64, vocab_offset=False):
    for dl_type, dl in [("train", train_ds), ("valid", valid_ds)]:
        print(f"Printing {dl_type} dataloader batch.")
        batch = next(iter(dl))
        print("Batch input_ids.shape", batch["input_ids"].shape)
        first = batch["input_ids"][0, :k]
        last = batch["input_ids"][0, -k:]
        print(f"First {k} tokens:", tokenizer.decode(first))
        print("ids:", first)
        print(f"Last {k} tokens:", tokenizer.decode(last))
        print("ids:", last)


def _generate_samples(diffusion_model, config, logger, tokenizer):
    logger.info("Starting Sample Eval.")
    if hasattr(config.data, "vocab_file"):
        vocab_offset = True 
        new_vocab = json.load(open(f'{config.data.cache_dir}/{config.data.vocab_file}', 'r'))
        tokenizer.add_tokens(new_vocab)
    model = _load_from_checkpoint(
        diffusion_model=diffusion_model, config=config, tokenizer=tokenizer
    )
    # model.metrics.gen_ppl.reset()
    # model.metrics.sample_entropy.reset()
    if config.eval.disable_ema:
        logger.info("Disabling EMA.")
        model.ema = None
    stride_length = config.sampling.stride_length
    num_strides = config.sampling.num_strides
    all_samples = [] 
    for _ in range(config.sampling.num_sample_batches):
        if config.sampling.semi_ar:
            _, intermediate_samples, _ = model.restore_model_and_semi_ar_sample(
                stride_length=stride_length,
                num_strides=num_strides,
                dt=1 / config.sampling.steps,
            )
            text_samples = intermediate_samples[-1]
            # Note: Samples generated using semi-ar method
            # need to to be processed before computing generative perplexity
            # since these samples contain numerous <|endoftext|> tokens
            # and diffusion.compute_generative_perplexity() discards
            # any text after the first EOS token.
        else:
            samples = model.restore_model_and_sample(num_steps=config.sampling.steps)

            model.metrics.record_entropy(samples)
            text_samples = model.tokenizer.batch_decode(samples)
            model.metrics.record_generative_perplexity(
                text_samples, config.model.length, model.device
            )
            all_samples.extend(list(text_samples))

    # pdb.set_trace()
    # with open('samples.txt', 'w') as f:
    #     for sample in text_samples:
    #         f.write(f"{sample}\n")

    # Evaluate text metrics
    # train_texts = open("/home/patrick/duo/train_texts.txt", "r").readlines()
    # text_metrics = evaluate_text_metrics(all_samples, train_texts)
    # print(text_metrics)
    if config.data.train == 'text8': 
        train_ds, valid_ds = dataloader.get_dataloaders(config, tokenizer)
        train_texts = []
        # check if train_texts.txt exists
        if os.path.exists("/home/patrick/duo/train_texts.txt"):
            print("Loading training texts from train_texts.txt for text metrics...")
            train_texts = open("/home/patrick/duo/train_texts.txt", "r").readlines()
        else: 
            print("Loading training texts for text metrics...")
            for batch in tqdm(train_ds):
                batch_texts = tokenizer.batch_decode(batch["input_ids"], skip_special_tokens=True)
                train_texts.extend(batch_texts)

            with open("/home/patrick/duo/train_texts.txt", "w") as f:
                for text in train_texts:
                    f.write(f"{text}\n")
        text_metrics = evaluate_text_metrics(all_samples, train_texts)
        
        print(text_metrics)
    else: 
        generative_ppl = 0.0
        entropy = 0.0
        if not config.sampling.semi_ar:
            generative_ppl = model.metrics.gen_ppl.compute().item()
            entropy = model.metrics.sample_entropy.compute().item()
            print("Generative perplexity:", generative_ppl)
            print("Sample entropy:", entropy)
        samples_path = config.eval.generated_samples_path
        with fsspec.open(samples_path, "w") as f:
            json.dump(
                {
                    "generative_ppl": generative_ppl,
                    "entropy": entropy,
                    "generated_seqs": all_samples,
                },
                f,
                indent=4,
            )
        print("Samples saved at:", samples_path)


def _eval_ppl(diffusion_model, config, logger, tokenizer):
    logger.info("Starting Perplexity Eval.")

    model = _load_from_checkpoint(
        diffusion_model=diffusion_model, config=config, tokenizer=tokenizer
    )
    if config.eval.disable_ema:
        logger.info("Disabling EMA.")
        model.ema = None

    wandb_logger = None
    if config.get("wandb", None) is not None:
        wandb_logger = L.pytorch.loggers.WandbLogger(
            config=omegaconf.OmegaConf.to_object(config), **config.wandb
        )
    callbacks = []
    if "callbacks" in config:
        for _, callback in config.callbacks.items():
            callbacks.append(hydra.utils.instantiate(callback))
    trainer = hydra.utils.instantiate(
        config.trainer,
        default_root_dir=os.getcwd(),
        callbacks=callbacks,
        strategy=hydra.utils.instantiate(config.strategy),
        logger=wandb_logger,
    )
    _, valid_ds = dataloader.get_dataloaders(
        config, tokenizer, skip_train=True, valid_seed=config.seed
    )
    trainer.validate(model, valid_ds)


def _train(diffusion_model, config, logger, tokenizer):
    logger.info("Starting Training.")
    wandb_logger = None
    if config.get("wandb", None) is not None:
        wandb_logger = L.pytorch.loggers.WandbLogger(
            config=omegaconf.OmegaConf.to_object(config), **config.wandb
        )

    if (
        config.checkpointing.resume_from_ckpt
        and config.checkpointing.resume_ckpt_path is not None
        and utils.fsspec_exists(config.checkpointing.resume_ckpt_path)
    ):
        ckpt_path = config.checkpointing.resume_ckpt_path
    else:
        ckpt_path = None

    # Lightning callbacks
    callbacks = []
    if "callbacks" in config:
        for _, callback in config.callbacks.items():
            callbacks.append(hydra.utils.instantiate(callback))

    train_ds, valid_ds = dataloader.get_dataloaders(config, tokenizer)

    if hasattr(config.data, "vocab_file"): 
        vocab_offset = True 
        new_vocab = json.load(open(f'{config.data.cache_dir}/{config.data.vocab_file}', 'r'))
        tokenizer.add_tokens(new_vocab)
    else: 
        vocab_offset = False
    _print_batch(train_ds, valid_ds, tokenizer, vocab_offset=vocab_offset)

    if config.training.finetune_path != "":
        assert utils.fsspec_exists(config.training.finetune_path)
        model = diffusion_model.load_from_checkpoint(
            config.training.finetune_path, tokenizer=tokenizer, config=config
        )
    else:
        model = diffusion_model(config, tokenizer=tokenizer)

    trainer = hydra.utils.instantiate(
        config.trainer,
        default_root_dir=os.getcwd(),
        callbacks=callbacks,
        strategy=hydra.utils.instantiate(config.strategy),
        logger=wandb_logger,
    )
    trainer.fit(model, train_ds, valid_ds, ckpt_path=ckpt_path)

@hydra.main(version_base=None, config_path="configs", config_name="config")
def main(config):
    """Main entry point for training."""
    L.seed_everything(config.seed)
    # _print_config(config, resolve=True, save_cfg=True)

    logger = utils.get_logger(__name__)
    tokenizer = dataloader.get_tokenizer(config)
    if config.algo.name == "ar":
        diffusion_model = algo.AR
    elif config.algo.name == "mdlm":
        diffusion_model = algo.MDLM
    elif config.algo.name == "duo_base":
        diffusion_model = algo.DUO_BASE
    elif config.algo.name == "d3pm":
        diffusion_model = algo.D3PMAbsorb
    elif config.algo.name == "sedd":
        diffusion_model = algo.SEDDAbsorb
    elif config.algo.name == "duo":
        diffusion_model = algo.DUO
    elif config.algo.name == "distillation":
        diffusion_model = algo.Distillation
    elif config.algo.name == "ot-finetune":
        diffusion_model = algo.OptimalTransportFinetune
    elif config.algo.name == "candi":
        diffusion_model = algo.CANDI
    else:
        raise ValueError(f"Invalid algorithm name: {config.algo.name}")
    kwargs = {
        "diffusion_model": diffusion_model,
        "config": config,
        "tokenizer": tokenizer,
        "logger": logger,
    }
    if config.mode == "sample_eval":
        _generate_samples(**kwargs)
    elif config.mode == "ppl_eval":
        _eval_ppl(**kwargs)
    else:
        _train(**kwargs)


if __name__ == "__main__":
    main()

	from tqdm import tqdm

	import json
	import os

	from typing import List, Dict, Tuple
	import fsspec
	import hydra
	import lightning as L
	import omegaconf
	import rich.syntax
	import rich.tree
	import torch
	import algo
	import dataloader as dataloader
	import utils
	from text_metrics import evaluate_text_metrics


	omegaconf.OmegaConf.register_new_resolver("cwd", os.getcwd)
	omegaconf.OmegaConf.register_new_resolver("device_count", torch.cuda.device_count)
	omegaconf.OmegaConf.register_new_resolver("eval", eval)
	omegaconf.OmegaConf.register_new_resolver("div_up", lambda x, y: (x + y - 1) // y)

	def _load_from_checkpoint(diffusion_model, config, tokenizer):
	if "hf" in config.algo.backbone:
	return diffusion_model(config, tokenizer=tokenizer).to("cuda")


	model = diffusion_model.load_from_checkpoint(
	config.eval.checkpoint_path, tokenizer=tokenizer, config=config
	)
	return model



	@L.pytorch.utilities.rank_zero_only
	def _print_config(
	config: omegaconf.DictConfig, resolve: bool = True, save_cfg: bool = True
	) -> None:
	"""Prints content of DictConfig using Rich library and its tree structure.

	Args:
	config (DictConfig): Configuration composed by Hydra.
	resolve (bool): Whether to resolve reference fields of DictConfig.
	save_cfg (bool): Whether to save the configuration tree to a file.
	"""

	style = "dim"
	tree = rich.tree.Tree("CONFIG", style=style, guide_style=style)

	fields = config.keys()
	for field in fields:
	branch = tree.add(field, style=style, guide_style=style)

	config_section = config.get(field)
	branch_content = str(config_section)
	if isinstance(config_section, omegaconf.DictConfig):
	branch_content = omegaconf.OmegaConf.to_yaml(
	config_section, resolve=resolve
	)

	branch.add(rich.syntax.Syntax(branch_content, "yaml"))
	rich.print(tree)
	if save_cfg:
	with fsspec.open(
	"{}/config_tree.txt".format(config.checkpointing.save_dir), "w"
	) as fp:
	rich.print(tree, file=fp)


	@L.pytorch.utilities.rank_zero_only
	def _print_batch(train_ds, valid_ds, tokenizer, k=64, vocab_offset=False):
	for dl_type, dl in [("train", train_ds), ("valid", valid_ds)]:
	print(f"Printing {dl_type} dataloader batch.")
	batch = next(iter(dl))
	print("Batch input_ids.shape", batch["input_ids"].shape)
	first = batch["input_ids"][0, :k]
	last = batch["input_ids"][0, -k:]
	print(f"First {k} tokens:", tokenizer.decode(first))
	print("ids:", first)
	print(f"Last {k} tokens:", tokenizer.decode(last))
	print("ids:", last)


	def _generate_samples(diffusion_model, config, logger, tokenizer):
	logger.info("Starting Sample Eval.")
	if hasattr(config.data, "vocab_file"):
	vocab_offset = True
	new_vocab = json.load(open(f'{config.data.cache_dir}/{config.data.vocab_file}', 'r'))
	tokenizer.add_tokens(new_vocab)
	model = _load_from_checkpoint(
	diffusion_model=diffusion_model, config=config, tokenizer=tokenizer
	)
	# model.metrics.gen_ppl.reset()
	# model.metrics.sample_entropy.reset()
	if config.eval.disable_ema:
	logger.info("Disabling EMA.")
	model.ema = None
	stride_length = config.sampling.stride_length
	num_strides = config.sampling.num_strides
	all_samples = []
	for _ in range(config.sampling.num_sample_batches):
	if config.sampling.semi_ar:
	_, intermediate_samples, _ = model.restore_model_and_semi_ar_sample(
	stride_length=stride_length,
	num_strides=num_strides,
	dt=1 / config.sampling.steps,
	)
	text_samples = intermediate_samples[-1]
	# Note: Samples generated using semi-ar method
	# need to to be processed before computing generative perplexity
	# since these samples contain numerous <\|endoftext\|> tokens
	# and diffusion.compute_generative_perplexity() discards
	# any text after the first EOS token.
	else:
	samples = model.restore_model_and_sample(num_steps=config.sampling.steps)

	model.metrics.record_entropy(samples)
	text_samples = model.tokenizer.batch_decode(samples)
	model.metrics.record_generative_perplexity(
	text_samples, config.model.length, model.device
	)
	all_samples.extend(list(text_samples))

	# pdb.set_trace()
	# with open('samples.txt', 'w') as f:
	# for sample in text_samples:
	# f.write(f"{sample}\n")

	# Evaluate text metrics
	# train_texts = open("/home/patrick/duo/train_texts.txt", "r").readlines()
	# text_metrics = evaluate_text_metrics(all_samples, train_texts)
	# print(text_metrics)
	if config.data.train == 'text8':
	train_ds, valid_ds = dataloader.get_dataloaders(config, tokenizer)
	train_texts = []
	# check if train_texts.txt exists
	if os.path.exists("/home/patrick/duo/train_texts.txt"):
	print("Loading training texts from train_texts.txt for text metrics...")
	train_texts = open("/home/patrick/duo/train_texts.txt", "r").readlines()
	else:
	print("Loading training texts for text metrics...")
	for batch in tqdm(train_ds):
	batch_texts = tokenizer.batch_decode(batch["input_ids"], skip_special_tokens=True)
	train_texts.extend(batch_texts)

	with open("/home/patrick/duo/train_texts.txt", "w") as f:
	for text in train_texts:
	f.write(f"{text}\n")
	text_metrics = evaluate_text_metrics(all_samples, train_texts)

	print(text_metrics)
	else:
	generative_ppl = 0.0
	entropy = 0.0
	if not config.sampling.semi_ar:
	generative_ppl = model.metrics.gen_ppl.compute().item()
	entropy = model.metrics.sample_entropy.compute().item()
	print("Generative perplexity:", generative_ppl)
	print("Sample entropy:", entropy)
	samples_path = config.eval.generated_samples_path
	with fsspec.open(samples_path, "w") as f:
	json.dump(
	{
	"generative_ppl": generative_ppl,
	"entropy": entropy,
	"generated_seqs": all_samples,
	},
	f,
	indent=4,
	)
	print("Samples saved at:", samples_path)


	def _eval_ppl(diffusion_model, config, logger, tokenizer):
	logger.info("Starting Perplexity Eval.")

	model = _load_from_checkpoint(
	diffusion_model=diffusion_model, config=config, tokenizer=tokenizer
	)
	if config.eval.disable_ema:
	logger.info("Disabling EMA.")
	model.ema = None

	wandb_logger = None
	if config.get("wandb", None) is not None:
	wandb_logger = L.pytorch.loggers.WandbLogger(
	config=omegaconf.OmegaConf.to_object(config), **config.wandb
	)
	callbacks = []
	if "callbacks" in config:
	for _, callback in config.callbacks.items():
	callbacks.append(hydra.utils.instantiate(callback))
	trainer = hydra.utils.instantiate(
	config.trainer,
	default_root_dir=os.getcwd(),
	callbacks=callbacks,
	strategy=hydra.utils.instantiate(config.strategy),
	logger=wandb_logger,
	)
	_, valid_ds = dataloader.get_dataloaders(
	config, tokenizer, skip_train=True, valid_seed=config.seed
	)
	trainer.validate(model, valid_ds)


	def _train(diffusion_model, config, logger, tokenizer):
	logger.info("Starting Training.")
	wandb_logger = None
	if config.get("wandb", None) is not None:
	wandb_logger = L.pytorch.loggers.WandbLogger(
	config=omegaconf.OmegaConf.to_object(config), **config.wandb
	)

	if (
	config.checkpointing.resume_from_ckpt
	and config.checkpointing.resume_ckpt_path is not None
	and utils.fsspec_exists(config.checkpointing.resume_ckpt_path)
	):
	ckpt_path = config.checkpointing.resume_ckpt_path
	else:
	ckpt_path = None

	# Lightning callbacks
	callbacks = []
	if "callbacks" in config:
	for _, callback in config.callbacks.items():
	callbacks.append(hydra.utils.instantiate(callback))

	train_ds, valid_ds = dataloader.get_dataloaders(config, tokenizer)

	if hasattr(config.data, "vocab_file"):
	vocab_offset = True
	new_vocab = json.load(open(f'{config.data.cache_dir}/{config.data.vocab_file}', 'r'))
	tokenizer.add_tokens(new_vocab)
	else:
	vocab_offset = False
	_print_batch(train_ds, valid_ds, tokenizer, vocab_offset=vocab_offset)

	if config.training.finetune_path != "":
	assert utils.fsspec_exists(config.training.finetune_path)
	model = diffusion_model.load_from_checkpoint(
	config.training.finetune_path, tokenizer=tokenizer, config=config
	)
	else:
	model = diffusion_model(config, tokenizer=tokenizer)

	trainer = hydra.utils.instantiate(
	config.trainer,
	default_root_dir=os.getcwd(),
	callbacks=callbacks,
	strategy=hydra.utils.instantiate(config.strategy),
	logger=wandb_logger,
	)
	trainer.fit(model, train_ds, valid_ds, ckpt_path=ckpt_path)

	@hydra.main(version_base=None, config_path="configs", config_name="config")
	def main(config):
	"""Main entry point for training."""
	L.seed_everything(config.seed)
	# _print_config(config, resolve=True, save_cfg=True)

	logger = utils.get_logger(__name__)
	tokenizer = dataloader.get_tokenizer(config)
	if config.algo.name == "ar":
	diffusion_model = algo.AR
	elif config.algo.name == "mdlm":
	diffusion_model = algo.MDLM
	elif config.algo.name == "duo_base":
	diffusion_model = algo.DUO_BASE
	elif config.algo.name == "d3pm":
	diffusion_model = algo.D3PMAbsorb
	elif config.algo.name == "sedd":
	diffusion_model = algo.SEDDAbsorb
	elif config.algo.name == "duo":
	diffusion_model = algo.DUO
	elif config.algo.name == "distillation":
	diffusion_model = algo.Distillation
	elif config.algo.name == "ot-finetune":
	diffusion_model = algo.OptimalTransportFinetune
	elif config.algo.name == "candi":
	diffusion_model = algo.CANDI
	else:
	raise ValueError(f"Invalid algorithm name: {config.algo.name}")
	kwargs = {
	"diffusion_model": diffusion_model,
	"config": config,
	"tokenizer": tokenizer,
	"logger": logger,
	}
	if config.mode == "sample_eval":
	_generate_samples(**kwargs)
	elif config.mode == "ppl_eval":
	_eval_ppl(**kwargs)
	else:
	_train(**kwargs)


	if __name__ == "__main__":
	main()