Initialize repository

.gitignore

+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+debug/
+logs/
+datasets/
+config/
+wandb/
+slurm/

data_utils.py

+"""Handle data processing and laoding."""
+
+from typing import Any, Optional, Tuple
+from pathlib import Path
+import json
+import logging
+import gzip
+import pickle as pkl
+
+import numpy as np
+from sklearn.utils import compute_class_weight
+from datasets import Dataset, DatasetDict, load_from_disk, disable_progress_bars
+
+
+_SEED = 42
+_DATA_DIR = (Path(__file__).parent / "datasets").resolve()
+_DATASETS_NAMES = [d.name for d in _DATA_DIR.iterdir() if d.is_dir()]
+
+
+def load_data(
+    dataset_name: str,
+    predefined_split_info: Optional[dict[str, Any]] = None,
+    use_validation: bool = True,
+    test_size: float = 0.15,
+) -> Tuple[Dataset, Dataset, Dataset]:
+    disable_progress_bars()
+    assert (
+        dataset_name in _DATASETS_NAMES
+    ), f"Dataset {dataset_name} not found. Must be one of {_DATASETS_NAMES}"
+    dataset_path = _DATA_DIR / dataset_name / "hf_dataset"
+    try:
+        if predefined_split_info:
+            logging.info(f"Loading dataset {dataset_name} from predefined splits")
+            train_data, test_data, val_data = _load_data_from_predefined_splits(
+                dataset_name, predefined_split_info, use_validation
+            )
+        else:
+            logging.info(f"Loading dataset {dataset_name} from HuggingFace dataset")
+            train_data, test_data, val_data = _load_data_from_hf_splits(
+                dataset_name, use_validation
+            )
+    except FileNotFoundError:
+        logging.warning(
+            f"Dataset {dataset_name} not found in {dataset_path}. Processing raw data."
+        )
+        train_data, test_data, val_data = preprocess_data(dataset_name, test_size, use_validation)
+        logging.info(f"Saving processed dataset to {dataset_path}")
+        if val_data is not None:
+            dataset_dict = DatasetDict(
+                {"train": train_data, "test": test_data, "validation": val_data}
+            )
+        else:
+            dataset_dict = DatasetDict({"train": train_data, "test": test_data})
+        dataset_dict.save_to_disk(dataset_path)
+
+    return train_data, test_data, val_data
+
+
+def _load_data_from_hf_splits(dataset_name, use_validation):
+    dataset_path = _DATA_DIR / dataset_name / "hf_dataset"
+    dataset_dict = load_from_disk(dataset_path)
+    train_data = dataset_dict["train"]
+    test_data = dataset_dict["test"]
+    if use_validation and "validation" not in dataset_dict:
+        raise FileNotFoundError
+    val_data = dataset_dict["validation"] if use_validation else None
+    return train_data, test_data, val_data
+
+
+def _load_data_from_predefined_splits(dataset_name, predefined_split_info, use_validation):
+    """Load dataset from pre-split txt files containing proteins IDs."""
+    dataset_dir = _DATA_DIR / dataset_name
+    assert dataset_dir.exists(), f"Dataset {dataset_name} not found in {dataset_dir}"
+    sequences = _unpickle(dataset_dir / "sequences.pkl")
+    labels = _unpickle(dataset_dir / "labels.pkl")
+
+    data = {"id": [], "sequence": [], "labels": []}
+    for key in labels.keys():
+        data["id"].append(key)
+        data["sequence"].append("".join(sequences[key]))
+        data["labels"].append(labels[key])
+        if len(sequences[key]) != len(labels[key]):
+            logging.warning(f"Sequence and label lengths do not match for {key}")
+    data = Dataset.from_dict(data)
+    splits_file = dataset_dir / predefined_split_info["splits_file"]  # relative path to JSON file
+    if not splits_file.exists():
+        raise FileNotFoundError(f"Splits dictionary not found in {splits_file}")
+
+    data_splits = {}
+    with open(splits_file, "r") as f:
+        splits_data = json.load(f)
+
+    for split_name in predefined_split_info["splits"]:
+        if "val" in split_name and not use_validation:
+            data_splits["val"] = None
+            continue
+        split_ids = splits_data[split_name]
+        data_splits[split_name] = data.filter(lambda x: x["id"] in split_ids)
+
+    return data_splits.values()  # data_splits["train"], data_splits["test"], data_splits["val"]
+
+
+def _unpickle(file: Path) -> Any:
+    try:
+        with open(file, "rb") as f:
+            return pkl.load(f)
+    except pkl.UnpicklingError:
+        with gzip.open(file, "rb") as f:
+            return pkl.load(f)
+
+
+def preprocess_data(
+    dataset_name: str, test_size=0.15, use_validation=True
+) -> Tuple[Dataset, Dataset]:
+
+    dataset_dir = _DATA_DIR / dataset_name
+    assert dataset_dir.exists(), f"Dataset {dataset_name} not found in {dataset_dir}"
+    sequences = _unpickle(dataset_dir / "sequences.pkl")
+    labels = _unpickle(dataset_dir / "labels.pkl")
+
+    data = {"id": [], "sequence": [], "labels": []}
+    for key in sequences.keys():
+        data["id"].append(key)
+        data["sequence"].append("".join(sequences[key]))
+        data["labels"].append(labels[key])
+        if len(sequences[key]) != len(labels[key]):
+            logging.warning(f"Sequence and label lengths do not match for {key}")
+    data = Dataset.from_dict(data)  # pd.DataFrame(data)
+
+    # split data into train, val with fixed seed
+    train_test_split = data.train_test_split(test_size=test_size, seed=_SEED)
+    train_data = train_test_split["train"]
+    test_data = train_test_split["test"]
+    val_data = None
+
+    if use_validation:
+        train_val_split = train_data.train_test_split(test_size=0.15, seed=_SEED)
+        train_data = train_val_split["train"]
+        val_data = train_val_split["test"]
+
+    return train_data, test_data, val_data
+
+
+def post_process_data(data: Dataset, max_length: int = None, pos_range=None) -> Dataset:
+    # cut sequences to max_length and labels accordingly
+    if max_length is not None:
+        data = data.map(
+            lambda x: {"sequence": x["sequence"][:max_length], "labels": x["labels"][:max_length]},
+            batched=False,
+        )
+    if pos_range is not None:
+        # remove samples with too few positive labels or too many positive labels
+        pos_frac = np.array([sum(labels) / len(labels) for labels in data["labels"]])
+        filter_mask = (pos_frac >= pos_range[0]) & (pos_frac <= pos_range[1])
+        data = data.select(np.where(filter_mask)[0])
+    return data
+
+
+def prune(data: Dataset, prune_frac: float = 0.2, random_seed: int = _SEED) -> Dataset:
+    # prune a fraction of the data
+    rng = np.random.default_rng(random_seed)
+    indices = rng.choice(len(data), int(len(data) * (1 - prune_frac)), replace=False)
+    return data.select(indices)
+
+
+def downsample_negatives(
+    X: np.ndarray, y: np.ndarray, downsample_ratio: float, random_seed: int = _SEED
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Downsample the vectors in X that have label 0, according to the specified downsample_ratio.
+
+    Args:
+    - X (np.ndarray): The input array of vectors (shape: [n_samples, n_features]).
+    - y (np.ndarray): The corresponding label array (shape: [n_samples, ]).
+    - downsample_ratio (float): The fraction of label 0 vectors to retain (0.0 < downsample_ratio <= 1.0).
+    - random_seed (int, optional): A random seed to ensure reproducibility. Defaults to None.
+
+    Returns:
+    - X_downsampled (np.ndarray): The downsampled array of vectors.
+    - y_downsampled (np.ndarray): The corresponding downsampled labels.
+    """
+    if downsample_ratio <= 0.0 or downsample_ratio > 1.0:
+        raise ValueError("downsample_ratio must be between 0.0 and 1.0")
+
+    # Set random seed for reproducibility
+    rng = np.random.default_rng(random_seed)
+
+    # Get the indices of vectors with label 0 and label 1
+    indices_negatives = np.where(y == 0)[0]
+    indices_positives = np.where(y == 1)[0]
+
+    num_to_keep = int(len(indices_negatives) * downsample_ratio)
+    indices_negatives_to_keep = rng.choice(indices_negatives, num_to_keep, replace=False)
+    indices_to_keep = np.concatenate([indices_negatives_to_keep, indices_positives])
+    np.random.shuffle(indices_to_keep)
+
+    # Return the downsampled X and y
+    return X[indices_to_keep], y[indices_to_keep]
+
+
+def get_class_weights(dataset: Dataset, label_column: str = "labels"):
+    labels = np.concatenate(dataset[label_column])
+    class_weights = compute_class_weight(
+        class_weight="balanced", classes=np.unique(labels), y=labels
+    )
+    return class_weights

hf_utils.py

+"""HugingFace models utils"""
+
+from pathlib import Path
+import logging
+import yaml
+
+import h5py
+import torch
+from transformers import (
+    AutoModelForTokenClassification,
+    AutoTokenizer,
+    DataCollatorForTokenClassification,
+    Trainer,
+    TrainingArguments,
+    EarlyStoppingCallback,
+    EvalPrediction,
+    set_seed as set_seed_hf,
+)
+from transformers.utils import logging as hf_logging
+from accelerate.utils import find_executable_batch_size
+import peft
+
+from utils import (
+    compute_metrics as _compute_metrics,
+    set_torch_device,
+    process_config_file,
+    HF_PRETRAINED_MODELS,
+)
+from torch_models import TorchModelHelper
+import data_utils
+
+
+_SEED = 42
+
+
+def set_seed():
+    set_seed_hf(_SEED)
+
+
+def set_verbosity(verbosity_level: int) -> None:
+    if verbosity_level == 0:
+        hf_logging.set_verbosity_error()
+    elif verbosity_level == 1:
+        hf_logging.set_verbosity_warning()
+    elif verbosity_level == 2:
+        hf_logging.set_verbosity_info()
+    else:
+        raise ValueError("Verbosity level must be between 0 and 2")
+
+
+def compute_metrics(eval_pred: EvalPrediction):
+    preds, labels = eval_pred
+    mask = labels != -100  # ignore padding tokens
+    labels = labels[mask]
+    preds = preds[mask]
+    return _compute_metrics(labels, preds)
+
+
+class HFModelHelper:
+    """Common class for loading `transformers.PreTrainedModel` models"""
+
+    def __init__(self, model_name: str, model_path: str = None, device: str = "cpu"):
+        model_id = HF_PRETRAINED_MODELS.get(model_name.upper(), model_name)
+        if model_path is None:
+            model_path = model_id
+        self.model_path = model_path
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_id, use_fast=("rostlab" not in model_id.lower())
+        )
+        # load base pre-trained model
+        try:
+            model = AutoModelForTokenClassification.from_pretrained(
+                model_path, num_labels=2, classifier_dropout=0.2
+            )
+            if model.dropout.p != 0.2:
+                logging.warning(
+                    "Classifier droput was not initialized to 0.2 Setting it manually."
+                )
+            model.dropout.p = 0.2
+            is_peft = False
+        except:
+            raise ValueError(f"Model {model_path} is not a valid Hugging Face model")
+        # load peft adapter if present
+        try:
+            model = peft.PeftModel.from_pretrained(model, model_id=model_path)
+            is_peft = True
+        except ValueError:
+            if Path(model_path).is_dir():
+                logging.warning(f"PEFT adapter not found at {model_path}. Loading base model.")
+
+        self.model = model
+        self.device = device
+        self.is_peft = is_peft
+
+        # Freeze all layers except the classifier for non-PEFT models
+        if not self.is_peft:
+            for name, param in self.model.named_parameters():
+                if "classifier" not in name:
+                    param.requires_grad = False
+
+    @property
+    def device(self):
+        return next(self.model.parameters()).device
+
+    @device.setter
+    def device(self, device: str):
+        self.model.to(device)
+
+    def get_model(self):
+        return self.model
+
+    def get_tokenizer(self):
+        return self.tokenizer
+
+    def get_collator(self):
+        return DataCollatorForTokenClassification(tokenizer=self.tokenizer)
+
+    def tokenize(self, sequences, **kwargs):
+        if "saprot" in self.tokenizer.name_or_path.lower():
+            # add mask token for structure vocabulary
+            if isinstance(sequences, str):
+                sequences = "#".join(sequences)
+            elif isinstance(sequences, list):
+                sequences = ["#".join(seq) for seq in sequences]
+        elif "rostlab" in self.tokenizer.name_or_path.lower():
+            if isinstance(sequences, str):
+                sequences = " ".join(sequences)
+            elif isinstance(sequences, list):
+                sequences = [" ".join(seq) for seq in sequences]
+        return self.tokenizer(sequences, **kwargs)
+
+    def convert_to_peft(self):
+        logging.info("Adding LoRA layers to model")
+        if self.is_peft:
+            logging.warning("Model is already a PEFT model. Skipping conversion.")
+            return
+        lora_config = peft.LoraConfig(
+            task_type=peft.TaskType.TOKEN_CLS,
+            r=16,
+            base_model_name_or_path=self.model_path,
+            target_modules=["q_proj", "v_proj", "query", "value", "q", "v"],
+            lora_dropout=0.2,
+            lora_alpha=32,
+        )
+        lora_model = peft.get_peft_model(self.model, lora_config)
+        lora_model.print_trainable_parameters()
+        self.model = lora_model
+        self.is_peft = True
+
+
+class BalancedTrainer(Trainer):
+    """Custom Trainer class for handling class imbalance in token classification tasks"""
+
+    def __init__(self, class_weights, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.class_weights = class_weights
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        """
+        How the loss is computed by Trainer. By default, all models return the loss in
+        the first element.
+        """
+        if self.label_smoother is not None and "labels" in inputs:
+            labels = inputs.pop("labels")
+        else:
+            labels = None
+        outputs = model(**inputs)
+
+        if labels is not None:
+            loss = self.label_smoother(outputs, labels)
+        else:
+
+            # Changes start here
+            # loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]
+            logits = outputs["logits"]
+            criterion = torch.nn.CrossEntropyLoss(weight=self.class_weights)
+            loss = criterion(logits.view(-1, 2), inputs["labels"].view(-1))
+            # Changes end here
+
+        return (loss, outputs) if return_outputs else loss
+
+
+def process_dataset(dataset, config, model_helper, split: str = "train"):
+    max_length = config["dataset"].get("max_length")
+    pos_range = config["dataset"].get("pos_range") if split == "train" else None
+
+    logging.info(f"Processing {split} data")
+    if max_length is not None:
+        max_length = max_length - 2  # account for special tokens
+    dataset = data_utils.post_process_data(dataset, max_length=max_length, pos_range=pos_range)
+
+    # prune train dataset if specified
+    prune_frac = config["dataset"].get("prune_frac")
+    if prune_frac is not None and split == "train":
+        logging.info(f"Pruning {prune_frac} of the data")
+        dataset = data_utils.prune(dataset, prune_frac=prune_frac)
+
+    def encode(sample):
+        return {
+            **model_helper.tokenize(
+                sample["sequence"], max_length=max_length, truncation=split != "test"
+            )
+        }
+
+    # tokenize sequences
+    logging.info(f"Tokenizing {split} data")
+    dataset = dataset.map(encode, batched=False)
+
+    dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])
+    logging.warning(f"Dataset processed: {split} samples: {len(dataset)}")
+    return dataset
+
+
+def get_class_weights(dataset, device="cpu"):
+    class_weights = data_utils.get_class_weights(dataset)
+    return torch.tensor(class_weights, device=device, dtype=torch.float)
+
+
+@find_executable_batch_size(starting_batch_size=64)
+def batch_finder(batch_size, model_helper, max_length):
+    """Simulate training step to find the maximum batch size that fits into GPU memory"""
+    assert (
+        torch.device(model_helper.device).type == "cuda"
+    ), "Batch finder only works with CUDA devices"
+    max_token_id = model_helper.get_tokenizer().vocab_size
+    if max_length is None:
+        max_length = 1024
+    dummy_input_ids = torch.randint(
+        0, max_token_id, (batch_size, max_length), device=model_helper.device
+    )
+    labels = torch.randint(0, 2, (batch_size, max_length), device=model_helper.device)
+    loss = model_helper.model(input_ids=dummy_input_ids, labels=labels).loss
+    loss.backward()
+
+    return batch_size // 2
+
+
+def load_data(config, model_helper):
+    # load and process data
+    predefined_split_info = config["dataset"].get("predefined_split_info", False)
+    train_data, test_data, val_data = data_utils.load_data(
+        dataset_name=config["dataset"]["name"],
+        predefined_split_info=predefined_split_info,
+        test_size=0.15,
+    )
+    train_data = process_dataset(train_data, config, model_helper, split="train")
+    val_data = process_dataset(val_data, config, model_helper, split="validation")
+    test_data = process_dataset(test_data, config, model_helper, split="test")
+    return train_data, val_data, test_data
+
+
+def get_model_helper(model_name, config):
+    device = set_torch_device()
+    if model_name.upper() in HF_PRETRAINED_MODELS:
+        logging.info(f"Loading Hugging Face model: {model_name}")
+        model_helper = HFModelHelper(model_name, device=device)
+        if config["model"]["is_peft"]:
+            model_helper.convert_to_peft()
+    else:
+        logging.info(f"Loading custom torch model: {model_name}")
+        model_config = config["model"]
+        model_config.pop("model_name", None)
+        try:
+            pretrained_model = model_config.pop("pretrained_model")
+        except KeyError:
+            raise ValueError(
+                "Pretrained model for downstrem torch model not found in configuration file"
+            )
+        model_helper = TorchModelHelper(model_name, pretrained_model, device, **model_config)
+    logging.debug(f"Model: {model_helper.model}")
+    return model_helper
+
+
+def get_training_arguments(train_config, model_helper, max_length, logging_dir):
+    train_config["logging_dir"] = logging_dir
+    train_config["output_dir"] = logging_dir
+
+    # # Warm-up run to initialize CUDA context
+    # logging.info("Performing warm-up run to initialize CUDA context")
+    # try:
+    #     batch_finder(batch_size=1, model_helper=model_helper, max_length=max_length)
+    # except RuntimeError as e:
+    #     logging.warning(f"Warm-up run failed: {e}")
+
+    # set batch size
+    logging.info("Finding maximum instantaneous batch size")
+    effective_batch_size = train_config.get("per_device_train_batch_size", 8)
+    max_train_batch_size = batch_finder(model_helper=model_helper, max_length=max_length)
+    train_batch_size = min(effective_batch_size, max_train_batch_size)
+    eval_batch_size = train_batch_size  # * 4
+    gradient_accumulation_steps = max(1, effective_batch_size // train_batch_size)
+    logging.info(
+        f"Batch size: Eval: {eval_batch_size}, train: {train_batch_size},"
+        f" gradient accumulation: {gradient_accumulation_steps}"
+    )
+    train_config["per_device_train_batch_size"] = train_batch_size
+    train_config["per_device_eval_batch_size"] = eval_batch_size
+    train_config["gradient_accumulation_steps"] = gradient_accumulation_steps
+    train_config["learning_rate"] = (
+        train_config.get("learning_rate", 5e-5) * gradient_accumulation_steps
+    )
+
+    train_args = TrainingArguments(**train_config)
+    return train_args
+
+
+def train(
+    model_name: str,
+    config_file: str,
+    logging_dir: str = "./log",
+    verbosity: int = 2,
+):
+    """Train function for Huggingface models"""
+
+    set_verbosity(verbosity)
+    set_seed()
+
+    # load config
+    config = process_config_file(config_file, logging_dir)
+    max_length = config["dataset"].get("max_length", 1024)
+
+    # load model and data
+    model_helper = get_model_helper(model_name, config)
+    train_data, val_data, test_data = load_data(config, model_helper)
+
+    # initialize Trainer
+    class_weights = get_class_weights(train_data, device=model_helper.device)
+    train_config = config["trainer"]
+    train_args = get_training_arguments(train_config, model_helper, max_length, logging_dir)
+    logging.info(f"Class weights: {class_weights}")
+    logging.debug(f"Training arguments: {train_args}")
+
+    trainer = BalancedTrainer(
+        class_weights=class_weights,
+        model=model_helper.model,
+        args=train_args,
+        data_collator=model_helper.get_collator(),
+        train_dataset=train_data,
+        eval_dataset=val_data,
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
+        compute_metrics=compute_metrics,
+    )
+
+    train_output = trainer.train()
+    train_metrics = trainer.metrics_format(train_output.metrics)
+    with open(f"{logging_dir}/train_metrics.yml", "a+") as file:
+        yaml.dump(train_metrics, file, default_flow_style=False)
+
+    # evaluate on test data
+    logging.info("Evaluating on test data")
+    test_metrics = trainer.evaluate(test_data)
+    with open(f"{logging_dir}/test_metrics.yml", "a+") as file:
+        yaml.dump(trainer.metrics_format(test_metrics), file, default_flow_style=False)
+    trainer.log_metrics("test", test_metrics)
+
+    best_model_checkpoint = trainer.state.best_model_checkpoint
+    return best_model_checkpoint
+
+
+def predict(checkpoint_dir, output_file: str, verbosity: int = 2):
+    """Predict function for Huggingface models"""
+
+    set_verbosity(verbosity)
+    device = set_torch_device()
+
+    test_data = 1
+    # load model
+    model_helper = HFModelHelper(checkpoint_dir, device)
+    args = TrainingArguments(per_device_eval_batch_size=1)
+
+    # For predictions, no need to use balanced trainer
+    trainer = Trainer(
+        model=model_helper.model,
+        args=args,
+        data_collator=model_helper.get_collator(),
+        compute_metrics=compute_metrics,
+    )
+
+    # compute predictions
+    preds = trainer.predict(test_data).predictions
+
+    # save predictions
+    Path(output_file).parent.mkdir(parents=True, exist_ok=True)
+    with h5py.File(output_file, "w") as f:
+        for i, sample in enumerate(test_data):
+            id = sample["id"]
+            f.create_dataset(id, data=preds[i])

hpo_optuna.py

+"""Hyperparameter search for HuggingFace models using Optuna backend."""
+
+from typing import Tuple, Callable, Any
+from pathlib import Path
+import os
+import sys
+import yaml
+import logging
+import importlib.util
+
+from transformers import AutoTokenizer, AutoModelForTokenClassification, EarlyStoppingCallback
+import peft
+
+sys.path.append(Path(__file__).parent)
+import utils
+import hf_utils
+
+
+class HPOModelHelper(hf_utils.HFModelHelper):
+    """
+    Provide access to hyperparameters initialization
+    and direct conversion to peft model.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        model_path: str = None,
+        device: str = "cpu",
+        dropout: float = 0.2,
+        lora_r: int = 16,
+        lora_target_modules: list = ["q_proj", "v_proj", "query", "value", "q", "v"],
+    ):
+        model_id = utils.HF_PRETRAINED_MODELS.get(model_name.upper(), model_name)
+        if model_path is None:
+            model_path = model_id
+        self.model_path = model_path
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_id, use_fast=("rostlab" not in model_id.lower())
+        )
+        # load base pre-trained model
+        model = AutoModelForTokenClassification.from_pretrained(
+            model_path, num_labels=2, classifier_dropout=dropout
+        )
+        if model.dropout.p != dropout:
+            logging.warning(
+                f"Classifier droput was not initialized to {dropout}. Setting it manually."
+            )
+        model.dropout.p = dropout
+
+        # load peft adapter if present
+        is_peft = False
+        try:
+            model = peft.PeftModel.from_pretrained(model, model_id=model_path)
+            is_peft = True
+        except ValueError:
+            logging.warning(
+                f"PEFT adapter not found at {model_path}. Loading base model and initializing lora"
+                " layers."
+            )
+
+        if not is_peft:
+            lora_config = peft.LoraConfig(
+                task_type=peft.TaskType.TOKEN_CLS,
+                r=lora_r,
+                base_model_name_or_path=self.model_path,
+                target_modules=lora_target_modules,
+                lora_dropout=dropout,
+                lora_alpha=2 * lora_r,
+            )
+            lora_model = peft.get_peft_model(model, lora_config)
+            lora_model.print_trainable_parameters()
+            self.model = lora_model
+        else:
+            self.model = model
+
+        self.device = device
+
+    def convert_to_peft(self):
+        raise NotImplementedError(
+            "Conversion to PEFT model is already done during initialization."
+        )
+
+
+def set_wandb_env(logging_dir: str):
+    if importlib.util.find_spec("wandb") is not None:
+        project = logging_dir.split("hflogs/")[-1].replace("/", "-")
+        os.environ["WANDB_PROJECT"] = project
+        os.environ["WANDB_WATCH"] = "parameters"
+        logging.info("wandb is available and environment variables are set.")
+    else:
+        logging.info("wandb is not available.")
+
+
+def _prepare_hpo(model_name: str, hpo_config: dict) -> Tuple[Callable[..., Any]]:
+    hyperparams = hpo_config["hyperparameters"]
+    objective = hpo_config.get("objective", "roc_auc")
+    logging.info(f"Objective metric: {objective}")
+
+    def optuna_hp_space(trial):
+        learning_rate = trial.suggest_float(
+            "learning_rate", **hyperparams["learning_rate"], log=True
+        )
+        weight_decay = trial.suggest_categorical("weight_decay", hyperparams["weight_decay"])
+        warmup_ratio = trial.suggest_categorical("warmup_ratio", hyperparams["warmup_ratio"])
+        params = {
+            "learning_rate": learning_rate,
+            "weight_decay": weight_decay,
+            "warmup_ratio": warmup_ratio,
+        }
+        logging.info(f"Trial hyperparameters: {params}")
+        return params
+
+    def compute_objective(metrics):
+        return metrics["eval_" + objective]
+
+    def model_init(trial):
+        """Initialize the model."""
+        # In the first initialization, trial is None and default values are used to avoid errors
+        # logging.info("Initializing model during trial")
+        dropout = (
+            trial.suggest_categorical("dropout", hyperparams["dropout"])
+            if "dropout" in hyperparams and trial is not None
+            else 0.2
+        )
+        lora_r = (
+            trial.suggest_categorical("lora_r", hyperparams["lora_r"])
+            if "lora_r" in hyperparams and trial is not None
+            else 16
+        )
+        lora_target_modules = (
+            trial.suggest_categorical("lora_target_modules", hyperparams["lora_target_modules"])
+            if "lora_target_modules" in hyperparams and trial is not None
+            else ["q_proj", "v_proj", "query", "value", "q", "v"]
+        )
+        # since optuna `suggest_categorical requires choices to be hashable, in the config file
+        # the choices are defined as strings separated by "__"
+        if isinstance(lora_target_modules, str):
+            lora_target_modules = lora_target_modules.split("__")
+        params = {"dropout": dropout, "lora_r": lora_r, "lora_target_modules": lora_target_modules}
+        if trial is not None:
+            logging.info(f"Model hyperparameters: {params}")
+        model_helper = HPOModelHelper(
+            model_name, **params
+        )  # may be necessary to add trial arguments
+        return model_helper.get_model()
+
+    return model_init, optuna_hp_space, compute_objective
+
+
+def run_hpo(
+    model_name: str,
+    config_file: str,
+    logging_dir: str = "./log",
+    verbosity: int = 2,
+):
+    hf_utils.set_seed()
+    hf_utils.set_verbosity(verbosity)
+    device = utils.set_torch_device()
+
+    # process config file and copy to logging directory
+    Path(logging_dir).mkdir(parents=True, exist_ok=True)
+    config = utils.process_config_file(config_file, logging_dir)
+    max_length = config["dataset"].get("max_length", 1024)
+
+    # initlialize base model helper for loading data and computing class weights
+    model_helper = hf_utils.HFModelHelper(model_name, device=device)
+    model_helper.convert_to_peft()
+    train_data, val_data, _ = hf_utils.load_data(config, model_helper)
+    print(train_data[0])
+    class_weights = hf_utils.get_class_weights(train_data, device=device)
+    data_collator = model_helper.get_collator()
+
+    # get default training arguments
+    set_wandb_env(logging_dir)
+    train_config = config["trainer"]
+    patience = train_config.pop("early_stopping_patience", None)
+    train_args = hf_utils.get_training_arguments(
+        train_config, model_helper, max_length, logging_dir
+    )
+    del model_helper  # remove model helper to free up memory
+
+    # get hpo util functions
+    hpo_config = config["hpo"]
+    model_init, optuna_hp_space, compute_objective = _prepare_hpo(model_name, hpo_config)
+
+    # initialize trainer
+    trainer = hf_utils.BalancedTrainer(
+        class_weights=class_weights,
+        model=None,
+        model_init=model_init,
+        args=train_args,
+        data_collator=data_collator,
+        train_dataset=train_data,
+        eval_dataset=val_data,
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=patience)] if patience else None,
+        compute_metrics=hf_utils.compute_metrics,
+    )
+
+    # run hyperparameter search
+    logging.info("Starting hyperparameter search...")
+    best_trial = trainer.hyperparameter_search(
+        direction="maximize",
+        hp_space=optuna_hp_space,
+        compute_objective=compute_objective,
+        n_trials=hpo_config.get("n_trials", 20),
+        backend="optuna",
+    )
+
+    # Save the best trial hyperparameters to a YAML file
+    result = dict()
+    result["run_id"] = best_trial.run_id
+    result["roc_auc"] = best_trial.objective
+    result["hyperparameters"] = best_trial.hyperparameters
+    output_file = Path(logging_dir) / "best_trial.yaml"
+    with open(output_file, "w") as file:
+        yaml.dump(result, file, default_flow_style=False)
+    logging.info(f"Best trial hyperparameters saved to {output_file}")
+
+
+if __name__ == "__main__":
+    args = utils.parse_cli(__doc__)
+    utils.setup_logging(args.verbosity)
+    logging.info(f"Starting hyperparameter search for model: {args.model_name}")
+    run_hpo(args.model_name, args.config_file, args.logging_dir, args.verbosity)

notebooks/caid_test.ipynb

+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Test LoRA-DR-suite models on CAID datasets\n",
+    "\n",
+    "Sample script to test ESM2_650M-LoRA-ID model on CAI3-Disorder_NOX dataset"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Import packages"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.nn.functional as F\n",
+    "\n",
+    "from datasets import Dataset\n",
+    "from sklearn import metrics\n",
+    "from tqdm.auto import tqdm\n",
+    "from transformers import AutoModelForTokenClassification, AutoTokenizer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "device(type='cuda')"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Set torch device\n",
+    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "DEVICE"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Download CAID data and convert them to HuggingFace Dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "vscode": {
+     "languageId": "shellscript"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
+      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
+      "100  143k  100  143k    0     0  1050k      0 --:--:-- --:--:-- --:--:-- 1057k\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Download fasta file with sequences and labels for the CAID challenge 3 and save it as \"caid3_disorder_nox.fasta\"\n",
+    "\n",
+    "! mkdir -p \"data\"\n",
+    "! curl https://caid.idpcentral.org/assets/sections/challenge/static/references/3/disorder_nox.fasta > \"data/caid3_disorder_nox.fasta\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Dataset({\n",
+       "    features: ['id', 'sequence', 'labels'],\n",
+       "    num_rows: 148\n",
+       "})"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Process FASTA file and convert data to Dataset\n",
+    "\n",
+    "# function to read in a fasta file and return a Hugging Face dataset\n",
+    "def parse_data(fasta_file):\n",
+    "    aas = set(\"ACDEFGHIKLMNPQRSTVWYXUZB\")\n",
+    "    data = {\"id\": [], \"sequence\": [], \"labels\": []}\n",
+    "    with open(fasta_file, \"r\") as f:\n",
+    "        lines = f.readlines()\n",
+    "        for i in range(0, len(lines), 3):\n",
+    "            assert lines[i].startswith(\">\"), f\"Line {i} is not a header\"\n",
+    "            data[\"id\"].append(lines[i].strip()[1:])\n",
+    "            sequence = lines[i+1].strip()\n",
+    "            assert set(sequence).issubset(aas), f\"Line {i+1} has invalid amino acids: {set(sequence)}\"\n",
+    "            data[\"sequence\"].append(lines[i+1].strip())\n",
+    "            labels = list(lines[i+2].strip())\n",
+    "            assert set(labels).issubset({\"0\", \"1\", \"-\"}), f\"Line {i+2} has invalid labels: {set(labels)}\"\n",
+    "            assert len(labels) == len(sequence), f\"Line {i+2} has different length than line {i+1}\"\n",
+    "            data[\"labels\"].append([int(x) if x != \"-\" else -100 for x in labels])       \n",
+    "    return Dataset.from_dict(data)\n",
+    "\n",
+    "test_dataset = parse_data(\"data/caid3_disorder_nox.fasta\")\n",
+    "test_dataset"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Tokenize data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Some weights of EsmForTokenClassification were not initialized from the model checkpoint at facebook/esm2_t33_650M_UR50D and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
+      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Load model and tokenizer\n",
+    "\n",
+    "model_id = \"CQSB/esm2_650M-LoRA-ID\"\n",
+    "model = AutoModelForTokenClassification.from_pretrained(model_id)\n",
+    "model.eval()\n",
+    "pretrained_model_id = model.config.name_or_path\n",
+    "tokenizer = AutoTokenizer.from_pretrained(pretrained_model_id)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "81322a2c1e50459699047ef094557c81",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Map:   0%|          | 0/148 [00:00<?, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "{'id': Value(dtype='string', id=None),\n",
+       " 'sequence': Value(dtype='string', id=None),\n",
+       " 'labels': Sequence(feature=Value(dtype='int64', id=None), length=-1, id=None),\n",
+       " 'input_ids': Sequence(feature=Value(dtype='int32', id=None), length=-1, id=None),\n",
+       " 'special_tokens_mask': Sequence(feature=Value(dtype='int8', id=None), length=-1, id=None),\n",
+       " 'attention_mask': Sequence(feature=Value(dtype='int8', id=None), length=-1, id=None)}"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "def tokenize_function(examples):\n",
+    "    return tokenizer(\n",
+    "        examples[\"sequence\"],\n",
+    "        padding=False,\n",
+    "        truncation=False,\n",
+    "        return_special_tokens_mask=True,\n",
+    "    )\n",
+    "\n",
+    "test_dataset = test_dataset.map(tokenize_function, batched=True)\n",
+    "test_dataset.features"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Compute predictions for input sequences"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@torch.inference_mode()\n",
+    "def compute_predictions(model, dataset):\n",
+    "    t0 = time.time()\n",
+    "    pred_dict = {}\n",
+    "    model = model.to(DEVICE)\n",
+    "    for sample in tqdm(dataset, desc=\"Computing predictions\", total=len(dataset)):\n",
+    "        input_ids = torch.tensor([sample[\"input_ids\"]], device=DEVICE)\n",
+    "        attention_mask = torch.tensor([sample[\"input_ids\"]], device=DEVICE)\n",
+    "        predictions = model(input_ids, attention_mask)[\"logits\"]\n",
+    "        predicted_proba = F.softmax(predictions, dim=-1).squeeze(0).cpu().numpy()\n",
+    "        predicted_proba = predicted_proba[np.array(sample[\"special_tokens_mask\"]) == 0, 1]\n",
+    "        pred_dict[sample[\"id\"]] = predicted_proba\n",
+    "    t1 = time.time() - t0\n",
+    "    print(f\"Total execution time: {t1:.3f} s | Average time per sequence: {t1/len(dataset):.3f}s\")\n",
+    "    return pred_dict\n",
+    "\n",
+    "predictions = compute_predictions(model, test_dataset)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Compute evaluation metrics "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Function for metrics computation\n",
+    "\n",
+    "eval_metrics = {\n",
+    "    \"f1\": metrics.f1_score,\n",
+    "    \"mcc\": metrics.matthews_corrcoef,\n",
+    "    \"balanced_accuracy\": metrics.balanced_accuracy_score,\n",
+    "    \"average_precision\": metrics.average_precision_score,\n",
+    "    \"precision\": metrics.precision_score,\n",
+    "    \"recall\": metrics.recall_score,\n",
+    "    \"roc_auc\": metrics.roc_auc_score,\n",
+    "    \"accuracy\": metrics.accuracy_score,\n",
+    "}\n",
+    "\n",
+    "\n",
+    "def compute_metrics(labels, preds):\n",
+    "    eval_results = {}\n",
+    "    for metric_name, metric_func in eval_metrics.items():\n",
+    "        try:\n",
+    "            eval_results[metric_name] = metric_func(labels, preds)\n",
+    "        except ValueError:\n",
+    "            try:\n",
+    "                eval_results[metric_name] = metric_func(labels, preds > 0.5)\n",
+    "            except Exception as e:\n",
+    "                raise e\n",
+    "    return eval_results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Correctly predicted sequences: 148/148\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Concatenate all predictions and labels\n",
+    "\n",
+    "all_labels = []\n",
+    "all_preds = []\n",
+    "labels = {sample[\"id\"]: np.array(sample[\"labels\"]) for sample in test_dataset}\n",
+    "\n",
+    "for id_, pred in predictions.items():\n",
+    "    all_labels.append(labels[id_])\n",
+    "    all_preds.append(pred)\n",
+    "print(f\"Correctly predicted sequences: {len(labels)}/{len(test_dataset)}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 38,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'f1': np.float64(0.6486735982897678),\n",
+       " 'mc': np.float64(0.5228634718918926),\n",
+       " 'balanced_accuracy': np.float64(0.7838498800558977),\n",
+       " 'average_precision': np.float64(0.7213002519436622),\n",
+       " 'precision': np.float64(0.5717773019271949),\n",
+       " 'recall': np.float64(0.7494667115751656),\n",
+       " 'roc_auc': np.float64(0.8799528631178382),\n",
+       " 'accuracy': 0.8014116225420191}"
+      ]
+     },
+     "execution_count": 38,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Compute metrics on labeled data\n",
+    "\n",
+    "all_labels = np.concatenate(all_labels)\n",
+    "all_preds = np.concatenate(all_preds)\n",
+    "mask = all_labels != -100\n",
+    "\n",
+    "eval_results = compute_metrics(all_labels[mask], all_preds[mask])\n",
+    "eval_results"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

requirements.txt

+# Installing PyTorch from a specific index URL, depending on CUDA version
+# In this case, CUDA 12.1
+--index-url https://download.pytorch.org/whl/cu121
+torch==2.3.0
+torchvision==0.18
+torchaudio==2.3.0
+
+# Return to PyPI for the rest of the packages
+--index-url https://pypi.org/simple
+h5py
+numpy==1.23.5
+scipy=1.15.0
+pandas==2.2.2
+scikit-learn==1.5.0
+transformers==4.45.0
+accelerate==1.0.0
+huggingface-hub==0.27.1
+peft==0.13.1
+datasets==3.0.1
+optuna==4.0.0
+wandb==0.18.5
+tqdm==4.67.1
+

utils.py

+"""Utils functions for setup, logging and training."""
+
+from pathlib import Path
+import logging
+import yaml
+import shutil
+
+from sklearn import metrics
+from scipy.special import softmax
+
+
+def setup_logging(verbosity: int) -> None:
+    if verbosity == 2:
+        level = logging.DEBUG
+    elif verbosity == 1:
+        level = logging.INFO
+    elif verbosity == 0:
+        level = logging.WARNING
+    logging.basicConfig(
+        level=level,
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+        handlers=[logging.StreamHandler()]
+    )
+    
+
+def parse_cli(description):
+    import argparse
+    
+    parser = argparse.ArgumentParser(
+        description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument(
+        "-m",
+        "--model-name",
+        type=str,
+        required=True,
+        help=(
+            "The name of the model to train. Must be a valid Hugging Face model, a local"
+            " checkpoint path, or 'sklearn'."
+        ),
+    )
+    parser.add_argument(
+        "-c",
+        "--config-file",
+        type=str,
+        required=True,
+        help="The path to the configuration file.",
+    )
+    parser.add_argument(
+        "-o",
+        "--logging-dir",
+        type=str,
+        default="./log",
+        help="The directory where logs and model checkpoints will be saved.",
+    )
+    parser.add_argument(
+        "--pred-file",
+        type=str,
+        default=None,
+        help="The path to the file where the predictions will be saved.",
+    )
+    parser.add_argument(
+        "-v",
+        "--verbosity",
+        type=int,
+        default=2,
+        help="The verbosity level of the logger.",
+    )
+    args = parser.parse_args()
+
+    # Check for invalid arguments and set defaults
+    if args.verbosity > 2:
+        logging.warning("Verbosity level must be between 0 and 2. Setting to 2.")
+        args.verbosity = 2
+    elif args.verbosity < 0:
+        logging.warning("Verbosity level must be between 0 and 2. Setting to 0.")
+        args.verbosity = 0
+
+    return args
+
+
+def process_config_file(config_file, logging_dir, sanity_check_fn=None):
+    # load config file
+    with open(config_file, "r") as file:
+        config = yaml.safe_load(file)
+    # run sanity check function
+    if sanity_check_fn:
+        sanity_check_fn(config)
+    # remove *.yml files from logging directory
+    for file in Path(logging_dir).glob("*.yml"):
+        file.unlink()
+    shutil.copy2(config_file, logging_dir)
+    return config
+
+
+# Pretrained models hub IDs
+HF_PRETRAINED_MODELS = {
+    "ESM2_35M": "facebook/esm2_t12_35M_UR50D",
+    "ESM2_650M": "facebook/esm2_t33_650M_UR50D",
+    "PROTT5": "Rostlab/prot_t5_xl_uniref50",
+    "ANKH_BASE": "ElnaggarLab/ankh-base",
+    "ANKH_LARGE": "ElnaggarLab/ankh-large",
+    "SAPROT_35M": "westlake-repl/SaProt_35M_AF2",
+    "SAPROT_650M": "westlake-repl/SaProt_650M_AF2",
+}
+
+
+_DEAFAULT_EVAL_METRICS = {
+    "f1": metrics.f1_score,
+    "mc": metrics.matthews_corrcoef,
+    "balanced_accuracy": metrics.balanced_accuracy_score,
+    "average_precision": metrics.average_precision_score,
+    "precision": metrics.precision_score,
+    "recall": metrics.recall_score,
+    "roc_auc": metrics.roc_auc_score,
+    "accuracy": metrics.accuracy_score,
+}
+
+
+def compute_metrics(labels, preds):
+    eval_results = {}
+    for metric_name, metric_func in _DEAFAULT_EVAL_METRICS.items():
+        try:
+            eval_results[metric_name] = metric_func(labels, softmax(preds, axis=-1)[:, 1])
+        except ValueError:
+            try:
+                eval_results[metric_name] = metric_func(labels, preds.argmax(axis=-1))
+            except Exception as e:
+                raise e
+    return eval_results
+
+
+def set_torch_device():
+    import torch
+
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available():
+        device = "mps"
+    else:
+        device = "cpu"
+    logging.info(f"Using default device: {device}")
+    return device