Wauplin HF staff commited on
Commit
0585716
1 Parent(s): 41cac13

better error message + fail early

Browse files
Files changed (2) hide show
  1. app.py +39 -43
  2. pyproject.toml +2 -0
app.py CHANGED
@@ -1,4 +1,6 @@
1
  import pathlib
 
 
2
  import tempfile
3
  from typing import Iterable, List
4
 
@@ -36,9 +38,7 @@ has_gpu = torch.cuda.is_available()
36
  # )
37
 
38
  cli = "mergekit-yaml config.yaml merge --copy-tokenizer" + (
39
- " --cuda --low-cpu-memory"
40
- if has_gpu
41
- else " --allow-crimes --out-shard-size 1B --lazy-unpickle"
42
  )
43
 
44
  ## This Space is heavily inspired by LazyMergeKit by Maxime Labonne
@@ -89,17 +89,22 @@ A quick overview of the currently supported merge methods:
89
  examples = [[str(f)] for f in pathlib.Path("examples").glob("*.yml")]
90
 
91
 
92
- def merge(
93
- yaml_config: str, hf_token: str | None, repo_name: str | None
94
- ) -> Iterable[List[Log]]:
95
  if not yaml_config:
96
- raise gr.Error("Empty yaml, pick an example below")
 
97
  try:
98
  merge_config = MergeConfiguration.model_validate(yaml.safe_load(yaml_config))
99
  except Exception as e:
100
- raise gr.Error(f"Invalid yaml {e}")
 
101
 
102
- runner = LogsViewRunner()
 
 
 
103
 
104
  with tempfile.TemporaryDirectory() as tmpdirname:
105
  tmpdir = pathlib.Path(tmpdirname)
@@ -109,44 +114,35 @@ def merge(
109
  config_path.write_text(yaml_config)
110
  yield runner.log(f"Merge configuration saved in {config_path}")
111
 
112
- if token is not None and repo_name == "":
113
- name = "-".join(
114
- model.model.path for model in merge_config.referenced_models()
115
- )
116
- repo_name = f"mergekit-{merge_config.merge_method}-{name}".replace(
117
- "/", "-"
118
- ).strip("-")
119
- if len(repo_name) > 50:
120
- repo_name = repo_name[:25] + "-etc-" + repo_name[25:]
121
- runner.log(f"Will save merged in {repo_name} once process is done.")
122
-
123
- if token is None:
124
- yield runner.log(
125
- "No token provided, merge will run in dry-run mode (no upload at the end of the process)."
126
- )
127
 
128
  yield from runner.run_command(cli.split(), cwd=merged_path)
129
 
130
  if runner.exit_code != 0:
131
- yield runner.log(
132
- "Merge failed. Terminating here. No model has been uploaded."
133
- )
134
  return
135
 
136
- if hf_token is not None:
137
-
138
- def upload_to_repo():
139
- api = huggingface_hub.HfApi(token=hf_token)
140
- print("Creating repo")
141
- repo_url = api.create_repo(repo_name, exist_ok=True)
142
- print(f"Repo created: {repo_url}")
143
- print("Starting upload")
144
- folder_url = api.upload_folder(
145
- repo_id=repo_url.repo_id, folder_path=merged_path / "merge"
146
- )
147
- print(f"Model successfully uploaded to {folder_url}")
148
-
149
- yield from runner.run_python(upload_to_repo)
150
 
151
 
152
  with gr.Blocks() as demo:
@@ -161,12 +157,12 @@ with gr.Blocks() as demo:
161
  label="HF Write Token",
162
  info="https://hf.co/settings/token",
163
  type="password",
164
- placeholder="optional, will not upload merge if empty (dry-run)",
165
  )
166
  repo_name = gr.Textbox(
167
  lines=1,
168
  label="Repo name",
169
- placeholder="optional, will create a random name if empty",
170
  )
171
  button = gr.Button("Merge", variant="primary")
172
  logs = LogsView()
 
1
  import pathlib
2
+ import random
3
+ import string
4
  import tempfile
5
  from typing import Iterable, List
6
 
 
38
  # )
39
 
40
  cli = "mergekit-yaml config.yaml merge --copy-tokenizer" + (
41
+ " --cuda --low-cpu-memory" if has_gpu else " --allow-crimes --out-shard-size 1B --lazy-unpickle"
 
 
42
  )
43
 
44
  ## This Space is heavily inspired by LazyMergeKit by Maxime Labonne
 
89
  examples = [[str(f)] for f in pathlib.Path("examples").glob("*.yml")]
90
 
91
 
92
+ def merge(yaml_config: str, hf_token: str, repo_name: str) -> Iterable[List[Log]]:
93
+ runner = LogsViewRunner()
94
+
95
  if not yaml_config:
96
+ yield runner.log("Empty yaml, pick an example below", level="ERROR")
97
+ return
98
  try:
99
  merge_config = MergeConfiguration.model_validate(yaml.safe_load(yaml_config))
100
  except Exception as e:
101
+ yield runner.log(f"Invalid yaml {e}", level="ERROR")
102
+ return
103
 
104
+ if not hf_token:
105
+ yield runner.log("You must provide a write-access token.", level="ERROR")
106
+ return
107
+ api = huggingface_hub.HfApi(token=hf_token)
108
 
109
  with tempfile.TemporaryDirectory() as tmpdirname:
110
  tmpdir = pathlib.Path(tmpdirname)
 
114
  config_path.write_text(yaml_config)
115
  yield runner.log(f"Merge configuration saved in {config_path}")
116
 
117
+ if not repo_name:
118
+ yield runner.log("No repo name provided. Generating a random one.")
119
+ repo_name = f"mergekit-{merge_config.merge_method}"
120
+ # Make repo_name "unique" (no need to be extra careful on uniqueness)
121
+ repo_name += "-" + "".join(random.choices(string.ascii_lowercase, k=7))
122
+ repo_name = repo_name.replace("/", "-").strip("-")
123
+
124
+ try:
125
+ yield runner.log(f"Creating repo {repo_name}")
126
+ repo_url = api.create_repo(repo_name, exist_ok=True)
127
+ yield runner.log(f"Repo created: {repo_url}")
128
+ except Exception as e:
129
+ yield runner.log(f"Error creating repo {e}", level="ERROR")
130
+ return
 
131
 
132
  yield from runner.run_command(cli.split(), cwd=merged_path)
133
 
134
  if runner.exit_code != 0:
135
+ yield runner.log("Merge failed. Deleting repo as no model is uploaded.", level="ERROR")
136
+ api.delete_repo(repo_url.repo_id)
 
137
  return
138
 
139
+ yield runner.log("Model merged successfully. Uploading to HF.")
140
+ yield from runner.run_python(
141
+ api.upload_folder,
142
+ repo_id=repo_url.repo_id,
143
+ folder_path=merged_path / "merge",
144
+ )
145
+ yield runner.log("Model successfully uploaded to HF.")
 
 
 
 
 
 
 
146
 
147
 
148
  with gr.Blocks() as demo:
 
157
  label="HF Write Token",
158
  info="https://hf.co/settings/token",
159
  type="password",
160
+ placeholder="Mandatory. Used to upload the merged model.",
161
  )
162
  repo_name = gr.Textbox(
163
  lines=1,
164
  label="Repo name",
165
+ placeholder="Optional. Will create a random name if empty.",
166
  )
167
  button = gr.Button("Merge", variant="primary")
168
  logs = LogsView()
pyproject.toml ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ [tool.ruff]
2
+ line-length=119