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()