Update app.py

app.py

@@ -1,96 +1,36 @@
-import panel as pn
-from transformers_agent_ui import TransformersAgentUI
-from huggingface_hub import HfApi
-
-# Load the model list from the script
-model_list = [
-    "LayoutLMv3",
-    "LayoutXLM",
-    "LED",
-    "LeViT",
-    "LiLT",
-    "LLaMA",
-    "Llama2",
-    "Llama3",
-    "LLaVa",
-    "LLaVA-NeXT",
-    "LLaVA-NeXT-Video",
-    "LLaVA-Onevision",
-    "Longformer",
-    "LongT5",
-    "LUKE",
-    "LXMERT",
-    "M-CTC-T",
-    "M2M100",
-    "MADLAD-400",
-    "Mamba",
-    "mamba2",
-    "Marian",
-    "MarkupLM",
-    "Mask2Former",
-    "MaskFormer",
-    "MatCha",
-    "mBART",
-    "mBART-50",
-    "MEGA",
-    "Megatron-BERT",
-    "Megatron-GPT2",
-    "MGP-STR",
-    "Mimi",
-    "Mistral",
-    "Mixtral",
-    "Mllama",
-    "mLUKE",
-    "MMS",
-    "MobileBERT",
-    "MobileNetV1",
-    "MobileNetV2",
-    "MobileViT",
-    "MobileViTV2",
-    "Moshi",
-    "MPNet",
-    "MPT",
-    "MRA",
-    "MT5",
-    "MusicGen",
-]
-
-# Create a dropdown menu for model selection
-model_selector = pn.widgets.Select(options=model_list, value=model_list[0])
-
-# Create a text input for the repository list file
-repo_list_file = pn.widgets.TextInput(value="repo_list.txt")
-
-# Create a text input for the output directory
-output_dir = pn.widgets.TextInput(value="output_dir")
-
-# Create a checkbox for dry run
-dry_run = pn.widgets.Checkbox(value=False)
-
-# Create a button to start the process
-start_button = pn.widgets.Button(name="Start")
-
-# Define the function to be executed when the button is clicked
-def start_process(event):
-    # Get the selected model and repository list file
-    model = model_selector.value
-    repo_list_file_path = repo_list_file.value
-    output_dir_path = output_dir.value
-    dry_run_value = dry_run.value
-
-    # Call the script with the selected options
-    process_repos(output_dir_path, model, repo_list_file_path, dry_run_value)
-
-# Link the button to the function
-start_button.on_click(start_process)
-
-# Create the UI
-ui = pn.Column(
-    pn.Row(model_selector, repo_list_file, output_dir, dry_run),
-    start_button,
+import gradio as gr
+import subprocess
+
+# Function to execute hf_merge.py with provided parameters
+def run_hf_merge(repo_list, output_dir, staging='./staging', p=0.5, lambda_val=1.0, dry_run=False):
+    command = [
+        "python3", "hf_merge.py", repo_list, output_dir,
+        "-staging", staging, "-p", str(p), "-lambda", str(lambda_val)
+    ]
+    if dry_run:
+        command.append("--dry")
+
+    result = subprocess.run(command, capture_output=True, text=True)
+    if result.returncode == 0:
+        return "Merge completed successfully.\n" + result.stdout
+    else:
+        return "Error during merge.\n" + result.stderr
+
+# Gradio interface
+interface = gr.Interface(
+    fn=run_hf_merge,
+    inputs=[
+        gr.Textbox(label="Repo List File"),
+        gr.Textbox(label="Output Directory"),
+        gr.Textbox(label="Staging Directory", value="./staging"),
+        gr.Number(label="Dropout Probability", value=0.5),
+        gr.Number(label="Scaling Factor", value=1.0),
+        gr.Checkbox(label="Dry Run")
+    ],
+    outputs="text",
+    title="HuggingFace Model Merger",
+    description="Merge HuggingFace models using the technique described in 'Language Models are Super Mario: Absorbing Abilities from Homologous Models as a Free Lunch'."
 )
 
-# Serve the UI
-if pn.state.served:
-    pn.extension("terminal", "notifications", notifications=True, design="bootstrap")
-    ui.servable()
+if __name__ == "__main__":
+    interface.launch()