omnigenome 0.3.0a0__py3-none-any.whl

omnigenome/auto/auto_bench/auto_bench.py

@@ -0,0 +1,484 @@
+# -*- coding: utf-8 -*-
+# file: auto_bench.py
+# time: 11:54 14/04/2024
+# author: YANG, HENG <hy345@exeter.ac.uk> (杨恒)
+# github: https://github.com/yangheng95
+# huggingface: https://huggingface.co/yangheng
+# google scholar: https://scholar.google.com/citations?user=NPq5a_0AAAAJ&hl=en
+# Copyright (C) 2019-2024. All Rights Reserved.
+
+import os
+import time
+import warnings
+
+import findfile
+import torch
+from metric_visualizer import MetricVisualizer
+
+from transformers import TrainingArguments, Trainer as HFTrainer
+from ...src.abc.abstract_tokenizer import OmniTokenizer
+from ...src.lora.lora_model import OmniLoraModel
+from ...src.misc.utils import (
+    seed_everything,
+    fprint,
+    load_module_from_path,
+    check_bench_version,
+    clean_temp_checkpoint,
+)
+from ...src.trainer.trainer import Trainer
+from ...src.trainer.accelerate_trainer import AccelerateTrainer
+from ...utility.hub_utils import download_benchmark
+from ... import __version__ as omnigenome_version
+
+
+class AutoBench:
+    """
+    AutoBench is a class for automatically benchmarking genomic foundation models.
+    
+    This class provides a comprehensive framework for evaluating genomic models
+    across multiple benchmarks and tasks. It handles loading benchmarks, models,
+    tokenizers, and running evaluations with proper metric tracking and result
+    visualization.
+    
+    AutoBench supports various evaluation scenarios including:
+    - Single model evaluation across multiple benchmarks
+    - Multi-seed evaluation for robustness testing
+    - Different trainer backends (native, accelerate, huggingface)
+    - Automatic metric visualization and result tracking
+    
+    Attributes:
+        benchmark (str): The name or path of the benchmark to use.
+        model_name_or_path (str): The name or path of the model to evaluate.
+        tokenizer: The tokenizer to use for evaluation.
+        autocast (str): The autocast precision to use ('fp16', 'bf16', etc.).
+        overwrite (bool): Whether to overwrite existing evaluation results.
+        trainer (str): The trainer to use ('native', 'accelerate', 'hf_trainer').
+        mv_path (str): Path to the metric visualizer file.
+        mv (MetricVisualizer): The metric visualizer instance.
+        bench_metadata: Metadata about the benchmark configuration.
+    """
+
+    def __init__(
+        self,
+        benchmark,
+        model_name_or_path,
+        tokenizer=None,
+        **kwargs,
+    ):
+        """
+        Initializes the AutoBench instance.
+
+        Args:
+            benchmark (str): The name or path of the benchmark to use.
+            model_name_or_path (str): The name or path of the model to evaluate.
+            tokenizer: The tokenizer to use. If None, it will be loaded from the model path.
+            **kwargs: Additional keyword arguments.
+                - autocast (str): The autocast precision to use ('fp16', 'bf16', etc.). 
+                  Defaults to 'fp16'.
+                - overwrite (bool): Whether to overwrite existing evaluation results. 
+                  Defaults to False.
+                - trainer (str): The trainer to use ('native', 'accelerate', 'hf_trainer'). 
+                  Defaults to 'native'.
+
+        Example:
+            >>> # Initialize with a benchmark and model
+            >>> bench = AutoBench("RGB", "model_name")
+            
+            >>> # Initialize with custom settings
+            >>> bench = AutoBench("RGB", "model_name", 
+            ...                   autocast="bf16", trainer="accelerate")
+        """
+        self.benchmark = benchmark.rstrip("/")
+        self.autocast = kwargs.pop("autocast", "fp16")
+        self.overwrite = kwargs.pop("overwrite", False)
+        self.trainer = kwargs.pop("trainer", "native")
+
+        self.model_name_or_path = model_name_or_path
+        self.tokenizer = tokenizer
+        if isinstance(self.model_name_or_path, str):
+            self.model_name_or_path = self.model_name_or_path.rstrip("/")
+            self.model_name = self.model_name_or_path.split("/")[-1]
+        else:
+            self.model_name = self.model_name_or_path.__class__.__name__
+        if isinstance(tokenizer, str):
+            self.tokenizer = tokenizer.rstrip("/")
+        os.makedirs("./autobench_evaluations", exist_ok=True)
+        time_str = time.strftime("%Y%m%d_%H%M%S", time.localtime())
+        mv_name = f"{benchmark}-{self.model_name}"
+        self.mv_path = f"./autobench_evaluations/{mv_name}-{time_str}.mv"
+
+        mv_paths = findfile.find_files(
+            "./autobench_evaluations",
+            and_key=[benchmark, self.model_name, ".mv"],
+        )
+        if mv_paths and not self.overwrite:
+            self.mv = MetricVisualizer.load(mv_paths[-1])
+            self.mv.summary(round=4)
+        else:
+            self.mv = MetricVisualizer(self.mv_path)
+        if not os.path.exists(self.benchmark):
+            fprint(
+                "Benchmark:",
+                benchmark,
+                "does not exist. Search online for available benchmarks.",
+            )
+            self.benchmark = download_benchmark(self.benchmark)
+
+        # Import benchmark list
+        self.bench_metadata = load_module_from_path(
+            f"bench_metadata", f"{self.benchmark}/metadata.py"
+        )
+        check_bench_version(
+            self.bench_metadata.__omnigenome_version__, omnigenome_version
+        )
+        fprint("Loaded benchmarks: ", self.bench_metadata.bench_list)
+        self.bench_info()
+
+    def bench_info(self):
+        """
+        Prints and returns information about the current benchmark setup.
+        
+        This method provides a comprehensive overview of the current
+        benchmark configuration, including benchmark details, model information,
+        and evaluation settings.
+
+        Returns:
+            str: A string containing benchmark information.
+
+        Example:
+            >>> info = bench.bench_info()
+            >>> print(info)
+        """
+        info = f"Benchmark Root: {self.benchmark}\n"
+        info += f"Benchmark List: {self.bench_metadata.bench_list}\n"
+        info += f"Model Name or Path: {self.model_name}\n"
+        info += f"Tokenizer: {self.tokenizer}\n"
+        info += f"Metric Visualizer Path: {self.mv_path}\n"
+        info += f"BenchConfig Details: {self.bench_metadata}\n"
+        fprint(info)
+        return info
+
+    def run(self, **kwargs):
+        """
+        Runs the benchmarking process.
+        
+        This method iterates through the tasks in the benchmark, loads the corresponding
+        configurations, initializes the model, tokenizer, and datasets, and then
+        trains and evaluates the model. It supports multiple evaluation seeds and
+        various trainer backends.
+
+        Args:
+            **kwargs: Additional keyword arguments that will override the default
+                     parameters in the benchmark configuration.
+
+        Example:
+            >>> # Run benchmarking with default settings
+            >>> bench.run()
+            
+            >>> # Run with custom parameters
+            >>> bench.run(learning_rate=1e-4, batch_size=16)
+        """
+        bs_scale = kwargs.pop("bs_scale", 1)
+        # Import benchmark config
+        for _, bench in enumerate(self.bench_metadata.bench_list):
+            clean_temp_checkpoint(1)  # clean temp checkpoint older than 1 day
+            fprint(
+                ">" * 80,
+                f"\nRunning evaluation for task: {bench}",
+                "Progress: ",
+                _ + 1,
+                "/",
+                len(self.bench_metadata.bench_list),
+                f"{(_ + 1) * 100 / len(self.bench_metadata.bench_list)}%",
+            )
+            bench_config_path = findfile.find_file(
+                self.benchmark, and_key=f"{self.benchmark}.{bench}.config".split(".")
+            )
+            config = load_module_from_path("config", bench_config_path)
+            bench_config = config.bench_config
+            fprint(f"Loaded config for {bench} from {bench_config_path}")
+            fprint(bench_config)
+
+            # Init Tokenizer and Model
+            if not self.tokenizer:
+                tokenizer = OmniTokenizer.from_pretrained(
+                    self.model_name_or_path,
+                    trust_remote_code=bench_config.get("trust_remote_code", True),
+                    **bench_config,
+                )
+            else:
+                tokenizer = self.tokenizer
+
+            if not isinstance(bench_config["seeds"], list):
+                bench_config["seeds"] = [bench_config["seeds"]]
+
+            random_seeds = bench_config["seeds"]
+            for seed in random_seeds:
+                _kwargs = kwargs.copy()
+                for key, value in _kwargs.items():
+                    if key in bench_config:
+                        fprint(
+                            "Override", key, "with", value, "according to the input kwargs"
+                        )
+                        bench_config.update({key: value})
+
+                    else:
+                        warnings.warn(
+                            f"kwarg: {key} not found in bench_config while setting {key} = {value}"
+                        )
+                        bench_config.update({key: value})
+
+                for key, value in bench_config.items():
+                    if key in bench_config and key in _kwargs:
+                        _kwargs.pop(key)
+
+                fprint(
+                    f"AutoBench Config for {bench}:",
+                    "\n".join([f"{k}: {v}" for k, v in bench_config.items()]),
+                )
+                for key, value in _kwargs.items():
+                    if key in bench_config:
+                        fprint(
+                            "Override", key, "with", value, "according to the input kwargs"
+                        )
+                        bench_config.update({key: value})
+
+                    else:
+                        warnings.warn(
+                            f"kwarg: {key} not found in bench_config while setting {key} = {value}"
+                        )
+                        bench_config.update({key: value})
+
+                for key, value in bench_config.items():
+                    if key in bench_config and key in _kwargs:
+                        _kwargs.pop(key)
+
+                fprint(
+                    f"AutoBench Config for {bench}:",
+                    "\n".join([f"{k}: {v}" for k, v in bench_config.items()]),
+                )
+
+                batch_size = (
+                    bench_config["batch_size"] if "batch_size" in bench_config else 8
+                ) * bs_scale
+
+                record_name = f"{self.benchmark}-{bench}-{self.model_name}".split("/")[
+                    -1
+                ]
+                # check if the record exists
+                if record_name in self.mv.transpose() and len(
+                    list(self.mv.transpose()[record_name].values())[0]
+                ) >= len(bench_config["seeds"]):
+                    continue
+
+                seed_everything(seed)
+                if self.model_name_or_path:
+                    model_cls = bench_config["model_cls"]
+                    model = model_cls(
+                        self.model_name_or_path,
+                        tokenizer=tokenizer,
+                        label2id=bench_config.label2id,
+                        num_labels=bench_config["num_labels"],
+                        trust_remote_code=True,
+                        ignore_mismatched_sizes=True,
+                    )
+                else:
+                    raise ValueError(
+                        "model_name_or_path is not specified. Please provide a valid model name or path."
+                    )
+
+                fprint(f"\n{model}")
+
+                if kwargs.get("lora_config", None) is not None:
+                    fprint("Applying LoRA to the model with config:", kwargs["lora_config"])
+                    model = OmniLoraModel(model, **kwargs.get("lora_config", {}))
+
+                # Init Trainer
+                dataset_cls = bench_config["dataset_cls"]
+
+                if hasattr(model.config, "max_position_embeddings"):
+                    max_length = min(
+                        bench_config["max_length"],
+                        model.config.max_position_embeddings,
+                    )
+                else:
+                    max_length = bench_config["max_length"]
+
+                train_set = dataset_cls(
+                    data_source=bench_config["train_file"],
+                    tokenizer=tokenizer,
+                    label2id=bench_config["label2id"],
+                    max_length=max_length,
+                    structure_in=bench_config.get("structure_in", False),
+                    max_examples=bench_config.get("max_examples", None),
+                    shuffle=bench_config.get("shuffle", True),
+                    drop_long_seq=bench_config.get("drop_long_seq", False),
+                    **_kwargs,
+                )
+                test_set = dataset_cls(
+                    data_source=bench_config["test_file"],
+                    tokenizer=tokenizer,
+                    label2id=bench_config["label2id"],
+                    max_length=max_length,
+                    structure_in=bench_config.get("structure_in", False),
+                    max_examples=bench_config.get("max_examples", None),
+                    shuffle=False,
+                    drop_long_seq=bench_config.get("drop_long_seq", False),
+                    **_kwargs,
+                )
+                valid_set = dataset_cls(
+                    data_source=bench_config["valid_file"],
+                    tokenizer=tokenizer,
+                    label2id=bench_config["label2id"],
+                    max_length=max_length,
+                    structure_in=bench_config.get("structure_in", False),
+                    max_examples=bench_config.get("max_examples", None),
+                    shuffle=False,
+                    drop_long_seq=bench_config.get("drop_long_seq", False),
+                    **_kwargs,
+                )
+
+                if self.trainer == "hf_trainer":
+                    # Set up HuggingFace Trainer
+                    hf_kwargs = {
+                        k: v
+                        for k, v in kwargs.items()
+                        if hasattr(TrainingArguments, k) and k != "output_dir"
+                    }
+                    training_args = TrainingArguments(
+                        output_dir=f"./autobench_evaluations/{self.model_name}-{bench}",
+                        num_train_epochs=hf_kwargs.pop(
+                            "num_train_epochs", bench_config["epochs"]
+                        ),
+                        per_device_train_batch_size=hf_kwargs.pop(
+                            "batch_size", batch_size
+                        ),
+                        per_device_eval_batch_size=hf_kwargs.pop(
+                            "batch_size", batch_size
+                        ),
+                        gradient_accumulation_steps=hf_kwargs.pop(
+                            "gradient_accumulation_steps", 1
+                        ),
+                        learning_rate=hf_kwargs.pop("learning_rate", 2e-5),
+                        weight_decay=hf_kwargs.pop("weight_decay", 0),
+                        eval_strategy=hf_kwargs.pop("eval_strategy", "epoch"),
+                        save_strategy=hf_kwargs.pop("save_strategy", "epoch"),
+                        fp16=hf_kwargs.pop("fp16", True),
+                        remove_unused_columns=False,
+                        label_names=["labels"],
+                        **hf_kwargs,
+                    )
+
+                    valid_set = valid_set if len(valid_set) else test_set
+
+                    if len(bench_config["compute_metrics"]) > 1:
+                        fprint(
+                            "Multiple metrics not supported by HFTrainer, using the first one metric only."
+                        )
+                    trainer = HFTrainer(
+                        model=model,
+                        args=training_args,
+                        train_dataset=train_set,
+                        eval_dataset=valid_set,
+                        compute_metrics=(
+                            bench_config["compute_metrics"][0]
+                            if isinstance(bench_config["compute_metrics"], list)
+                            else bench_config["compute_metrics"]
+                        ),
+                    )
+
+                    # Train and evaluate
+                    eval_result = trainer.evaluate(
+                        valid_set if len(valid_set) else test_set
+                    )
+                    print(eval_result)
+                    train_result = trainer.train()
+                    eval_result = trainer.evaluate()
+                    test_result = trainer.evaluate(
+                        test_set if len(test_set) else valid_set
+                    )
+
+                    metrics = {
+                        "train": train_result.metrics,
+                        "eval": eval_result,
+                        "test": test_result,
+                    }
+                    fprint(metrics)
+                else:
+                    optimizer = torch.optim.AdamW(
+                        filter(lambda p: p.requires_grad, model.parameters()),
+                        lr=(
+                            bench_config["learning_rate"]
+                            if "learning_rate" in bench_config
+                            else 2e-5
+                        ),
+                        weight_decay=(
+                            bench_config["weight_decay"]
+                            if "weight_decay" in bench_config
+                            else 0
+                        ),
+                    )
+                    if self.trainer == "accelerate":
+                        trainer_cls = AccelerateTrainer
+                    else:
+                        trainer_cls = Trainer
+                    fprint(f"Using Trainer: {trainer_cls}")
+                    trainer = trainer_cls(
+                        model=model,
+                        train_dataset=train_set,
+                        eval_dataset=valid_set,
+                        test_dataset=test_set,
+                        batch_size=batch_size,
+                        patience=(
+                            bench_config["patience"]
+                            if "patience" in bench_config
+                            else 3
+                        ),
+                        epochs=bench_config["epochs"],
+                        gradient_accumulation_steps=bench_config.get(
+                            "gradient_accumulation_steps", 1
+                        ),
+                        optimizer=optimizer,
+                        loss_fn=(
+                            bench_config["loss_fn"]
+                            if "loss_fn" in bench_config
+                            else None
+                        ),
+                        compute_metrics=bench_config["compute_metrics"],
+                        seed=seed,
+                        autocast=self.autocast,
+                        **_kwargs,
+                    )
+                    metrics = trainer.train()
+
+                    predictions = trainer.predictions
+
+                    if bench_config.get("save_predictions", False):
+                        os.makedirs(f"predictions/{bench}", exist_ok=True)
+                        import numpy as np
+
+                        for split in predictions.keys():
+                            with open(
+                                f"predictions/{bench}/{split}.npy",
+                                "wb",
+                            ) as f:
+                                np.save(f, predictions[split])
+
+                    if metrics:
+                        for key, value in metrics["test"][-1].items():
+                            try:
+                                value = float(value)
+                            except:
+                                pass  # ignore non-float values
+                            self.mv.log(f"{record_name}", f"{key}", value)
+                        # for key, value in metrics['test'][-1].items():
+                        #     self.mv.log(f'{record_name}', f'test_{key}', value)
+                        # for i, valid_metrics in enumerate(metrics["valid"]):
+                        #     for key, value in valid_metrics.items():
+                        #         self.mv.log(f'{record_name}', f'valid_epoch_{i}_{key}', value)
+
+                        self.mv.summary(round=4)
+                        self.mv.dump(self.mv_path)
+                        self.mv.to_csv(self.mv_path.replace(".mv", ".csv"))
+                    del model, trainer, optimizer
+                    torch.cuda.empty_cache()