app.py

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'."
)

if __name__ == "__main__":
    interface.launch()