|
import pathlib |
|
import subprocess |
|
import tempfile |
|
from typing import Generator |
|
import gradio as gr |
|
import huggingface_hub |
|
import torch |
|
import yaml |
|
|
|
has_gpu = torch.cuda.is_available() |
|
|
|
cli = "mergekit-yaml config.yaml merge --copy-tokenizer" + ( |
|
" --cuda --low-cpu-memory" |
|
if has_gpu |
|
else " --allow-crimes --out-shard-size 1B --lazy-unpickle" |
|
) |
|
|
|
print(cli) |
|
|
|
|
|
|
|
|
|
|
|
MARKDOWN_DESCRIPTION = """ |
|
# mergekit-gui |
|
|
|
The fastest way to perform a model merge π₯ |
|
|
|
Specify a YAML configuration file (see examples below) and a HF token and this app will perform the merge and upload the merged model to your user profile. |
|
""" |
|
|
|
MARKDOWN_ARTICLE = """ |
|
___ |
|
|
|
## Merge Configuration |
|
|
|
[Mergekit](https://github.com/arcee-ai/mergekit) configurations are YAML documents specifying the operations to perform in order to produce your merged model. |
|
Below are the primary elements of a configuration file: |
|
|
|
- `merge_method`: Specifies the method to use for merging models. See [Merge Methods](https://github.com/arcee-ai/mergekit#merge-methods) for a list. |
|
- `slices`: Defines slices of layers from different models to be used. This field is mutually exclusive with `models`. |
|
- `models`: Defines entire models to be used for merging. This field is mutually exclusive with `slices`. |
|
- `base_model`: Specifies the base model used in some merging methods. |
|
- `parameters`: Holds various parameters such as weights and densities, which can also be specified at different levels of the configuration. |
|
- `dtype`: Specifies the data type used for the merging operation. |
|
- `tokenizer_source`: Determines how to construct a tokenizer for the merged model. |
|
|
|
## Merge Methods |
|
|
|
A quick overview of the currently supported merge methods: |
|
|
|
| Method | `merge_method` value | Multi-Model | Uses base model | |
|
| -------------------------------------------------------------------------------------------- | -------------------- | ----------- | --------------- | |
|
| Linear ([Model Soups](https://arxiv.org/abs/2203.05482)) | `linear` | β
| β | |
|
| SLERP | `slerp` | β | β
| |
|
| [Task Arithmetic](https://arxiv.org/abs/2212.04089) | `task_arithmetic` | β
| β
| |
|
| [TIES](https://arxiv.org/abs/2306.01708) | `ties` | β
| β
| |
|
| [DARE](https://arxiv.org/abs/2311.03099) [TIES](https://arxiv.org/abs/2306.01708) | `dare_ties` | β
| β
| |
|
| [DARE](https://arxiv.org/abs/2311.03099) [Task Arithmetic](https://arxiv.org/abs/2212.04089) | `dare_linear` | β
| β
| |
|
| Passthrough | `passthrough` | β | β | |
|
| [Model Stock](https://arxiv.org/abs/2403.19522) | `model_stock` | β
| β
| |
|
|
|
""" |
|
|
|
examples = [[f.name, f.read_text()] for f in pathlib.Path("examples").glob("*.yml")] |
|
|
|
|
|
def merge( |
|
example_filename: str, yaml_config: str, hf_token: str, repo_name: str |
|
) -> Generator[str, None, None]: |
|
output = "" |
|
if not yaml_config: |
|
raise gr.Error("Empty yaml, pick an example below") |
|
try: |
|
_ = yaml.safe_load(yaml_config) |
|
except: |
|
raise gr.Error("Invalid yaml") |
|
|
|
with tempfile.TemporaryDirectory() as tmpdirname: |
|
tmpdir = pathlib.Path(tmpdirname) |
|
output += f"About to start merging in directory {tmpdir}\n\n" |
|
yield output |
|
with open(tmpdir / "config.yaml", "w", encoding="utf-8") as f: |
|
f.write(yaml_config) |
|
output += cli + "\n\n" |
|
yield output |
|
|
|
cmd = cli.split() |
|
popen = subprocess.Popen( |
|
cmd, |
|
cwd=tmpdir, |
|
stdout=subprocess.PIPE, |
|
stderr=subprocess.STDOUT, |
|
universal_newlines=True, |
|
) |
|
for stdout_line in iter(popen.stdout.readline, ""): |
|
output += stdout_line |
|
yield output |
|
popen.stdout.close() |
|
return_code = popen.wait() |
|
if return_code: |
|
raise gr.Error(f"subprocess error: {return_code}") |
|
|
|
|
|
demo = gr.Interface( |
|
description=MARKDOWN_DESCRIPTION, |
|
article=MARKDOWN_ARTICLE, |
|
fn=merge, |
|
inputs=[ |
|
gr.Textbox(visible=False, label="filename"), |
|
gr.Code( |
|
language="yaml", |
|
lines=10, |
|
label="config.yaml", |
|
), |
|
gr.Textbox( |
|
lines=1, |
|
label="HF Write Token", |
|
info="https://hf.co/settings/token", |
|
type="password", |
|
placeholder="optional, will not upload merge if empty (dry-run)", |
|
), |
|
gr.Textbox( |
|
lines=1, |
|
label="Repo name", |
|
placeholder="optional, will create a random name if empty", |
|
), |
|
], |
|
outputs=gr.Textbox(label="output", lines=12, show_copy_button=True), |
|
allow_flagging="never", |
|
submit_btn="Merge", |
|
examples=examples, |
|
cache_examples=False, |
|
).queue(default_concurrency_limit=1) |
|
|
|
|
|
demo.launch() |
|
|