diff --git a/src/autotrain/cli/autotrain.py b/src/autotrain/cli/autotrain.py index 59c80a445f..25b016f57d 100644 --- a/src/autotrain/cli/autotrain.py +++ b/src/autotrain/cli/autotrain.py @@ -12,6 +12,7 @@ from .run_tabular import RunAutoTrainTabularCommand from .run_text_classification import RunAutoTrainTextClassificationCommand from .run_token_classification import RunAutoTrainTokenClassificationCommand +from .run_tools import RunAutoTrainToolsCommand def main(): @@ -35,6 +36,7 @@ def main(): RunAutoTrainSpaceRunnerCommand.register_subcommand(commands_parser) RunAutoTrainSeq2SeqCommand.register_subcommand(commands_parser) RunAutoTrainTokenClassificationCommand.register_subcommand(commands_parser) + RunAutoTrainToolsCommand.register_subcommand(commands_parser) args = parser.parse_args() diff --git a/src/autotrain/cli/run_tools.py b/src/autotrain/cli/run_tools.py new file mode 100644 index 0000000000..c8fe1369f1 --- /dev/null +++ b/src/autotrain/cli/run_tools.py @@ -0,0 +1,99 @@ +from argparse import ArgumentParser + +from . import BaseAutoTrainCommand + + +def run_tools_command_factory(args): + return RunAutoTrainToolsCommand(args) + + +class RunAutoTrainToolsCommand(BaseAutoTrainCommand): + @staticmethod + def register_subcommand(parser: ArgumentParser): + run_app_parser = parser.add_parser("tools", help="Run AutoTrain tools") + subparsers = run_app_parser.add_subparsers(title="tools", dest="tool_name") + + merge_llm_parser = subparsers.add_parser( + "merge-llm-adapter", + help="Merge LLM Adapter tool", + ) + merge_llm_parser.add_argument( + "--base-model-path", + type=str, + help="Base model path", + ) + merge_llm_parser.add_argument( + "--adapter-path", + type=str, + help="Adapter path", + ) + merge_llm_parser.add_argument( + "--token", + type=str, + help="Token", + default=None, + required=False, + ) + merge_llm_parser.add_argument( + "--pad-to-multiple-of", + type=int, + help="Pad to multiple of", + default=None, + required=False, + ) + merge_llm_parser.add_argument( + "--output-folder", + type=str, + help="Output folder", + required=False, + default=None, + ) + merge_llm_parser.add_argument( + "--push-to-hub", + action="store_true", + help="Push to Hugging Face Hub", + required=False, + ) + merge_llm_parser.set_defaults(func=run_tools_command_factory, merge_llm_adapter=True) + + convert_to_kohya_parser = subparsers.add_parser("convert_to_kohya", help="Convert to Kohya tool") + convert_to_kohya_parser.add_argument( + "--input-path", + type=str, + help="Input path", + ) + convert_to_kohya_parser.add_argument( + "--output-path", + type=str, + help="Output path", + ) + convert_to_kohya_parser.set_defaults(func=run_tools_command_factory, convert_to_kohya=True) + + def __init__(self, args): + self.args = args + + def run(self): + if getattr(self.args, "merge_llm_adapter", False): + self.run_merge_llm_adapter() + if getattr(self.args, "convert_to_kohya", False): + self.run_convert_to_kohya() + + def run_merge_llm_adapter(self): + from autotrain.tools.merge_adapter import merge_llm_adapter + + merge_llm_adapter( + base_model_path=self.args.base_model_path, + adapter_path=self.args.adapter_path, + token=self.args.token, + output_folder=self.args.output_folder, + pad_to_multiple_of=self.args.pad_to_multiple_of, + push_to_hub=self.args.push_to_hub, + ) + + def run_convert_to_kohya(self): + from autotrain.tools.convert_to_kohya import convert_to_kohya + + convert_to_kohya( + input_path=self.args.input_path, + output_path=self.args.output_path, + ) diff --git a/src/autotrain/templates/login.html b/src/autotrain/templates/login.html index 2da5230d90..4e7e3630de 100644 --- a/src/autotrain/templates/login.html +++ b/src/autotrain/templates/login.html @@ -7,6 +7,13 @@ + +
@@ -15,15 +22,23 @@
-

Login

-

Please login to use +

Please login to use AutoTrain

-
- + +

Alternatively, if you face login issues, you can add your + Hugging Face Write Token to this space as a secret in space settings. Note: The name of secret must be + HF_TOKEN and the value must be your Hugging Face WRITE token! You can find your tokens in user settings.

+
diff --git a/src/autotrain/tests/test_cli.py b/src/autotrain/tests/test_cli.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/autotrain/tools/__init__.py b/src/autotrain/tools/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/autotrain/tools/convert_to_kohya.py b/src/autotrain/tools/convert_to_kohya.py new file mode 100644 index 0000000000..c42ab117e5 --- /dev/null +++ b/src/autotrain/tools/convert_to_kohya.py @@ -0,0 +1,13 @@ +from diffusers.utils import convert_all_state_dict_to_peft, convert_state_dict_to_kohya +from safetensors.torch import load_file, save_file + +from autotrain import logger + + +def convert_to_kohya(input_path, output_path): + logger.info(f"Converting Lora state dict from {input_path} to Kohya state dict at {output_path}") + lora_state_dict = load_file(input_path) + peft_state_dict = convert_all_state_dict_to_peft(lora_state_dict) + kohya_state_dict = convert_state_dict_to_kohya(peft_state_dict) + save_file(kohya_state_dict, output_path) + logger.info(f"Kohya state dict saved at {output_path}") diff --git a/src/autotrain/tools/merge_adapter.py b/src/autotrain/tools/merge_adapter.py new file mode 100644 index 0000000000..2b3df5a502 --- /dev/null +++ b/src/autotrain/tools/merge_adapter.py @@ -0,0 +1,52 @@ +import torch +from peft import PeftModel +from transformers import AutoModelForCausalLM, AutoTokenizer + +from autotrain import logger +from autotrain.trainers.common import ALLOW_REMOTE_CODE + + +def merge_llm_adapter( + base_model_path, adapter_path, token, output_folder=None, pad_to_multiple_of=None, push_to_hub=False +): + + if output_folder is None and push_to_hub is False: + raise ValueError("You must specify either --output_folder or --push_to_hub") + + logger.info("Loading adapter...") + base_model = AutoModelForCausalLM.from_pretrained( + base_model_path, + torch_dtype=torch.float16, + low_cpu_mem_usage=True, + trust_remote_code=ALLOW_REMOTE_CODE, + token=token, + ) + + tokenizer = AutoTokenizer.from_pretrained( + adapter_path, + trust_remote_code=ALLOW_REMOTE_CODE, + token=token, + ) + if pad_to_multiple_of: + base_model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=pad_to_multiple_of) + else: + base_model.resize_token_embeddings(len(tokenizer)) + + model = PeftModel.from_pretrained( + base_model, + adapter_path, + token=token, + ) + model = model.merge_and_unload() + + if output_folder is not None: + logger.info("Saving target model...") + model.save_pretrained(output_folder) + tokenizer.save_pretrained(output_folder) + logger.info(f"Model saved to {output_folder}") + + if push_to_hub: + logger.info("Pushing model to Hugging Face Hub...") + model.push_to_hub(adapter_path) + tokenizer.push_to_hub(adapter_path) + logger.info(f"Model pushed to Hugging Face Hub as {adapter_path}")