Update app.py

app.py

@@ -27,100 +27,71 @@ 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`        | ✅          | ✅              |
-## Citation
-This GUI is powered by [Arcee's MergeKit](https://arxiv.org/abs/2403.13257).
-If you use it in your research, please cite the following paper:
-```
-@article{goddard2024arcee,
-  title={Arcee's MergeKit: A Toolkit for Merging Large Language Models},
-  author={Goddard, Charles and Siriwardhana, Shamane and Ehghaghi, Malikeh and Meyers, Luke and Karpukhin, Vlad and Benedict, Brian and McQuade, Mark and Solawetz, Jacob},
-  journal={arXiv preprint arXiv:2403.13257},
-  year={2024}
-}
-```
-This Space is heavily inspired by LazyMergeKit by Maxime Labonne (see [Colab](https://colab.research.google.com/drive/1obulZ1ROXHjYLn6PPZJwRR6GzgQogxxb)).
-"""
-
 examples = [[str(f)] for f in pathlib.Path("examples").glob("*.yaml")]
 COMMUNITY_HF_TOKEN = os.getenv("COMMUNITY_HF_TOKEN")
 
-def merge(yaml_config: str, hf_token: str, repo_name: str, profile_name: str) -> Iterable[List[Log]]:
+def merge_multiple_methods(yaml_config: str, hf_token: str, repo_name: str, profile_name: str) -> Iterable[List[Log]]:
     runner = LogsViewRunner()
 
     if not yaml_config:
         yield runner.log("Empty yaml, pick an example below", level="ERROR")
         return
+
     try:
         merge_config = MergeConfiguration.model_validate(yaml.safe_load(yaml_config))
     except Exception as e:
         yield runner.log(f"Invalid yaml {e}", level="ERROR")
         return
 
-    is_community_model = False
-    if not hf_token:
-        if "/" in repo_name and not repo_name.startswith("mergekit-community/"):
-            yield runner.log(
-                f"Cannot upload merge model to namespace {repo_name.split('/')[0]}: you must provide a valid token.",
-                level="ERROR",
-            )
-            return
-        yield runner.log(
-            "No HF token provided. Your merged model will be uploaded to the https://huggingface.co/mergekit-community organization."
-        )
-        is_community_model = True
-        if not COMMUNITY_HF_TOKEN:
-            raise gr.Error("Cannot upload to community org: community token not set by Space owner.")
-        hf_token = COMMUNITY_HF_TOKEN
+    methods_to_merge = ['dare_ties', 'slerp', 'ties']
+    current_yaml_config = yaml_config
+    merged_model_path = None
 
-    api = huggingface_hub.HfApi(token=hf_token)
+    for method in methods_to_merge:
+        yield from run_merge_for_method(method, current_yaml_config, hf_token, repo_name, profile_name, runner)
+        current_yaml_config = get_merged_yaml(current_yaml_config, method)
+        yield runner.log(f"Model merged with {method}. Proceeding to next method...")
+
+    merged_model_path = "final_merged_model"  # Placeholder, adjust based on your process
+
+    if merged_model_path:
+        yield runner.log(f"Model successfully merged using all methods. Saving unified model to {merged_model_path}")
+        # Save final YAML
+        example_yaml = generate_example_yaml(methods_to_merge)
+        yield runner.log(f"Generated example YAML: {example_yaml}")
+
+        # Here, you could potentially upload the final merged model
+        # Upload logic goes here if needed
+
+def get_merged_yaml(original_yaml: str, method: str) -> str:
+    yaml_data = yaml.safe_load(original_yaml)
+    yaml_data['merge_method'] = method
+    return yaml.dump(yaml_data)
+
+def run_merge_for_method(method: str, yaml_config: str, hf_token: str, repo_name: str, profile_name: str, runner: LogsViewRunner):
+    yaml_data = yaml.safe_load(yaml_config)
+    yaml_data['merge_method'] = method
+    new_yaml_config = yaml.dump(yaml_data)
 
     with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname:
         tmpdir = pathlib.Path(tmpdirname)
         merged_path = tmpdir / "merged"
         merged_path.mkdir(parents=True, exist_ok=True)
         config_path = merged_path / "config.yaml"
-        config_path.write_text(yaml_config)
-        yield runner.log(f"Merge configuration saved in {config_path}")
+        config_path.write_text(new_yaml_config)
+        yield runner.log(f"Merge configuration saved for {method} in {config_path}")
 
         if not repo_name:
-            yield runner.log("No repo name provided. Generating a random one.")
-            repo_name = f"{profile_name}/mergekit-{merge_config.merge_method}" if profile_name else f"mergekit-{merge_config.merge_method}"
+            repo_name = f"{profile_name}/mergekit-{method}" if profile_name else f"mergekit-{method}"
             repo_name += "-" + "".join(random.choices(string.ascii_lowercase, k=7))
             repo_name = repo_name.replace("/", "-").strip("-")
 
-        if is_community_model and not repo_name.startswith("mergekit-community/"):
-            repo_name = f"mergekit-community/{repo_name}"
-
         try:
-            yield runner.log(f"Creating repo {repo_name}")
-            repo_url = api.create_repo(repo_name, exist_ok=True)
-            yield runner.log(f"Repo created: {repo_url}")
+            yield runner.log(f"Creating repo for {method} {repo_name}")
+            repo_url = huggingface_hub.HfApi(token=hf_token).create_repo(repo_name, exist_ok=True)
+            yield runner.log(f"Repo created for {method}: {repo_url}")
         except Exception as e:
-            yield runner.log(f"Error creating repo {e}", level="ERROR")
+            yield runner.log(f"Error creating repo for {method}: {e}", level="ERROR")
             return
 
         tmp_env = os.environ.copy()
@@ -129,17 +100,49 @@ def merge(yaml_config: str, hf_token: str, repo_name: str, profile_name: str) ->
         yield from runner.run_command(full_cli.split(), cwd=merged_path, env=tmp_env)
 
         if runner.exit_code != 0:
-            yield runner.log("Merge failed. Deleting repo as no model is uploaded.", level="ERROR")
-            api.delete_repo(repo_url.repo_id)
+            yield runner.log(f"Merge for {method} failed. Deleting repo as no model is uploaded.", level="ERROR")
+            huggingface_hub.HfApi(token=hf_token).delete_repo(repo_url.repo_id)
             return
 
-        yield runner.log("Model merged successfully. Uploading to HF.")
+        yield runner.log(f"Model merged with {method}. Uploading to HF.")
         yield from runner.run_python(
-            api.upload_folder,
+            huggingface_hub.HfApi(token=hf_token).upload_folder,
             repo_id=repo_url.repo_id,
             folder_path=merged_path / "merge",
         )
-        yield runner.log(f"Model successfully uploaded to HF: {repo_url.repo_id}")
+        yield runner.log(f"Model successfully uploaded to HF with {method}: {repo_url.repo_id}")
+
+def generate_example_yaml(methods: List[str]) -> str:
+    """Genera un archivo YAML de ejemplo que refleja la secuencia de métodos de fusión aplicados"""
+    example_yaml = {
+        'merge_method': 'linear',  # O el método final que decidas usar
+        'models': ['model1', 'model2', 'model3'],  # Ejemplo de modelos a fusionar
+        'slices': None,  # Puedes agregar slices si es necesario
+        'parameters': {
+            'normalize': False,
+            'weight': 0.5
+        },
+        'tokenizer_source': 'union',  # Definir el tokenizer
+    }
+    # Agregar los métodos de fusión aplicados
+    example_yaml['merge_method_sequence'] = methods
+
+    return yaml.dump(example_yaml)
+
+def merge(yaml_config: str, hf_token: str, repo_name: str, profile_name: str) -> Iterable[List[Log]]:
+    runner = LogsViewRunner()
+
+    if not yaml_config:
+        yield runner.log("Empty yaml, pick an example below", level="ERROR")
+        return
+
+    try:
+        merge_config = MergeConfiguration.model_validate(yaml.safe_load(yaml_config))
+    except Exception as e:
+        yield runner.log(f"Invalid yaml {e}", level="ERROR")
+        return
+
+    yield from merge_multiple_methods(yaml_config, hf_token, repo_name, profile_name)
 
 with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN_DESCRIPTION)
@@ -190,4 +193,4 @@ def _garbage_collect_every_hour():
 pool = ThreadPoolExecutor()
 pool.submit(_garbage_collect_every_hour)
 
-demo.queue(default_concurrency_limit=2).launch()
+demo.queue(default_concurrency_limit=2).launch()