doc-to-lora/hyperlora/utils.py

import ast
import os
import random
import string
import time
import yaml
import logging
from contextlib import contextmanager
from typing import Iterable, Optional


import torch
from peft import PeftConfig, PeftModel
from peft.tuners.tuners_utils import BaseTunerLayer, check_target_module_exists
from peft.utils import get_peft_model_state_dict

logger = logging.getLogger()


# taken from https://discuss.pytorch.org/t/opinion-eval-should-be-a-context-manager/18998/3
@contextmanager
def evaluating(*models):
    """Temporarily switch to evaluation mode."""
    is_training = [model.training if model is not None else False for model in models]
    try:
        for model in models:
            if model is not None:
                model.eval()
        yield models
    finally:
        for model, training in zip(models, is_training):
            if model is not None:
                model.train(training)


def get_layers(model):
    if hasattr(model, "model"):
        return get_layers(model.model)
    return model.layers


def get_num_layers(model):
    return len(get_layers(model))


def get_num_params(model):
    total_params = 0
    trainable_params = 0
    for p in model.parameters():
        total_params += p.numel()
        if p.requires_grad:
            trainable_params += p.numel()

    return total_params, trainable_params


def log_num_train_params(model):
    logger.debug("Trainable model parameters:")
    for name, p in model.named_parameters():
        if p.requires_grad:
            logger.debug(f"{name}, dtype:{p.dtype}")

    num_total_params, num_trainable_params = get_num_params(model)
    logger.info(
        f"trainable params: {num_trainable_params:,d} "
        f"|| all params: {num_total_params:,d} "
        f"|| trainable%: {100 * num_trainable_params / num_total_params:.4f}"
    )


def get_run_name():
    uuid = "".join(
        [random.choice(string.ascii_letters + string.digits) for _ in range(8)]
    )
    run_name = time.strftime("%Y%m%d-%H%M%S") + f"_{uuid}"
    return run_name


def try_convert(s):
    try:
        return ast.literal_eval(s)
    except:
        return s


def extract_cli_args(argv: list[str]):
    out = dict()
    for elem in argv:
        if elem.endswith(".yaml"):
            out["config"] = elem

        elif elem.startswith("--"):
            k, v = elem.split("=")
            k = k.split("--")[1]
            v = try_convert(v)
            # if k.startswith('env_'):
            #     k = k.split('_')[1]
            out[k] = v
    return out


def setup_logging(output_dir, debug=False):
    global logger

    os.makedirs(output_dir, exist_ok=True)

    log_formatter = logging.Formatter(
        fmt="%(asctime)s %(levelname)s: %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
    )
    stream_level = logging.DEBUG if debug else logging.INFO
    stream_handler = logging.StreamHandler()
    stream_handler.setFormatter(log_formatter)
    stream_handler.setLevel(stream_level)
    logger.addHandler(stream_handler)

    log_path = f"{output_dir}/debug.log"
    debug_handler = logging.FileHandler(log_path, delay=True)
    debug_handler.setLevel(logging.DEBUG)
    debug_handler.setFormatter(log_formatter)
    logger.addHandler(debug_handler)
    logger.setLevel(logging.DEBUG)
    logger.info(f"Logging to: {log_path}")


def validate_columns(tokenized_ds):
    cols = ["input_ids", "attention_mask", "labels"]
    if "ctx_features" in tokenized_ds["train"].column_names:
        cols += ["ctx_features", "ctx_attn_mask"]
    ref_cols = set(cols)
    assert (
        set(tokenized_ds["train"].column_names) == ref_cols
    ), f"Columns mismatch: {set(tokenized_ds['train'].column_names)} != {ref_cols}"


def validate_args(args_list):
    # there shouldn't be overlap between args
    keys = set()
    for args in args_list:
        args_keys = set(vars(args).keys())
        assert len(keys & args_keys) == 0, "Overlap between args"
        keys |= args_keys


def save_yaml(data, path):
    # Filter out non-primitive fields
    data = {
        k: v
        for k, v in data.items()
        if isinstance(v, (int, float, str, bool, list, dict, type(None)))
    }

    with open(path, "w") as file:
        yaml.dump(data, file)


def get_peft_in_out_features(
    model: PeftModel,
    peft_config: Optional[PeftConfig] = None,
) -> tuple[dict[str, int], dict[str, int]]:

    if peft_config is None:
        peft_config = model.peft_config["default"]
    in_features = dict()
    out_features = dict()
    for module_name, module in model.named_modules():
        if not check_target_module_exists(peft_config, module_name):
            continue
        if not isinstance(module, BaseTunerLayer):
            continue
        # support just Linear layer for now
        # all modules should be a leave module that is Linear layer
        assert isinstance(
            module.base_layer, torch.nn.Linear
        ), "all modules should be a leave module that is Linear layer"

        # this should always pass
        name = module_name.split(".")[-1]
        assert name in peft_config.target_modules

        if name not in in_features:
            in_features[name] = module.in_features
            out_features[name] = module.out_features
        else:
            # assumes each module has the same input and output features
            assert in_features[name] == module.in_features
            assert out_features[name] == module.out_features

    return in_features, out_features


def generated_lora_to_state_dict(
    lora_dict: dict,
    module_names: dict,
    target_modules: list[str],
    layer_indices: Iterable[int],
) -> dict:
    lora_state_dict = dict()
    for target_module in target_modules:
        for layer_idx in layer_indices:
            for module_name in module_names[target_module][layer_idx]:
                if "lora_A" in module_name:
                    lora_state_dict[module_name] = (
                        lora_dict[target_module]["A"][layer_idx].cpu().contiguous()
                    )
                elif "lora_B" in module_name:
                    lora_state_dict[module_name] = (
                        lora_dict[target_module]["B"][layer_idx].cpu().contiguous()
                    )
                else:
                    raise ValueError(f"Unexpected module name: {module_name}")
    return lora_state_dict


def get_lora_module_names(
    model: PeftModel,
    target_modules: list[str],
    layer_indices: Iterable[int],
) -> dict[str, list[str]]:
    module_names = {
        target_module: [[] for _ in range(len(layer_indices))]
        for target_module in target_modules
    }
    for k in get_peft_model_state_dict(model):
        if "lora" not in k:
            continue
        layer_idx = int(k.split("layers.")[-1].split(".")[0])
        if layer_idx in layer_indices:
            for target_module in target_modules:
                if target_module in k:
                    module_names[target_module][layer_idx].append(k)
                    break
    return module_names