Merge pull request #32 from alex-jw-brooks/merge_models

Merge models

Author: hickeyma

tuning/utils/merge_model_utils.py

@@ -0,0 +1,91 @@
+# Standard
+from typing import Union
+import argparse
+import json
+import os
+
+# Third Party
+from peft import PeftModel
+from tqdm import tqdm
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+def create_merged_model(
+    checkpoint_models: Union[str, list[str]],
+    export_path: str = None,
+    base_model: str = None,
+    save_tokenizer: bool = True,
+):
+    """Given a base model & checkpoint model(s) which were tuned with lora, load into memory
+    & create a merged model. If an export path is specified, write it to disk. If multiple
+    checkpoint models are provided, we merge_and_unload() them one after the other, which
+    combines them with equal weights.
+
+    TODO: In the future, it's probably a good idea to explore different combination schemes,
+    which can likely be done using a combination of add_weighted_adapter() and merge_and_unload().
+
+    References:
+    - https://github.com/huggingface/peft/issues/1040
+    - https://github.com/huggingface/peft/issues/280#issuecomment-1500805831
+    - https://huggingface.co/docs/peft/main/en/package_reference/lora#peft.LoraModel.add_weighted_adapter
+
+    Args:
+        checkpoint_model: Union[str, list[str]]
+            One or more lora checkpoints containing adapters.
+        export_path: str
+            Path to export the merged model to.
+        base_model: str
+            Base model to be leveraged. If no base model is specified, the base model is pulled
+            from the checkpoint model's adapter config.
+        save_tokenizer: bool
+            Indicates whether or not we should save the tokenizer from the base model. Only
+            used if the export_path is set.
+
+    Returns:
+        transformers model
+            Merged model created from the checkpoint / base model.
+    """
+    if isinstance(checkpoint_models, str):
+        checkpoint_models = [checkpoint_models]
+
+    if base_model is None:
+        base_model = fetch_base_model_from_checkpoint(checkpoint_models)
+    model = AutoModelForCausalLM.from_pretrained(base_model)
+
+    # Merge each of the lora adapter models into the base model with equal weights
+    for checkpoint_model in tqdm(checkpoint_models):
+        model = PeftModel.from_pretrained(model, checkpoint_model)
+        model = model.merge_and_unload()
+
+    if export_path is not None:
+        model.save_pretrained(export_path)
+        # Export the tokenizer into the merged model dir
+        if save_tokenizer:
+            tokenizer = AutoTokenizer.from_pretrained(base_model)
+            tokenizer.save_pretrained(export_path)
+    return model
+
+
+def fetch_base_model_from_checkpoint(checkpoint_model: str) -> str:
+    """Inspects the checkpoint model, locates the adapter config, and grabs the
+    base_model_name_or_path.
+
+    Args:
+        checkpoint_model: str
+            Checkpoint model containing the adapter config, which specifies the base model.
+
+    Returns:
+        str
+            base_model_name_or_path specified in the adapter config of the tuned peft model.
+    """
+    adapter_config = os.path.join(checkpoint_model, "adapter_config.json")
+    if not os.path.isfile(adapter_config):
+        raise FileNotFoundError("Unable to locate adapter config to infer base model!")
+
+    with open(adapter_config, "r") as cfg:
+        adapter_dict = json.load(cfg)
+    if "base_model_name_or_path" not in adapter_dict:
+        raise KeyError(
+            "Base model adapter config exists, but has no base_model_name_or_path!"
+        )
+    return adapter_dict["base_model_name_or_path"]