Update app.py

app.py

@@ -1,42 +1,35 @@
 import gradio as gr
 import subprocess
-import spaces
 
-@spaces.GPU
-# Function to execute hf_merge.py with provided parameters
-def run_hf_merge(repo_names, output_dir, staging='./staging', p=0.5, lambda_val=1.0, dry_run=False):
-    repo_list = "\n".join(repo_names.split(","))
-    with open("repo_list.txt", "w") as file:
-        file.write(repo_list)
-    
+def merge_models(weight_drop_prob, scaling_factor, base_model, model_to_merge, output_path):
+    # Construct the command to run hf_merge.py
     command = [
-        "python3", "hf_merge.py", "repo_list.txt", output_dir,
-        "-staging", staging, "-p", str(p), "-lambda", str(lambda_val)
+        "python3", "hf_merge.py",
+        "-p", str(weight_drop_prob),
+        "-lambda", str(scaling_factor),
+        base_model, model_to_merge, output_path
     ]
-    if dry_run:
-        command.append("--dry")
-
+    
+    # Run the command and capture the output
     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
+    
+    # Return the output of the command
+    return result.stdout
 
-# Gradio interface
-interface = gr.Interface(
-    fn=run_hf_merge,
+# Define the Gradio interface
+iface = gr.Interface(
+    fn=merge_models,
     inputs=[
-        gr.Textbox(label="Repository Names (comma-separated)"),
-        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")
+        gr.inputs.Slider(minimum=0, maximum=1, default=0.13, label="Weight Drop Probability"),
+        gr.inputs.Number(default=3.0, label="Scaling Factor"),
+        gr.inputs.Textbox(label="Base Model File/Folder"),
+        gr.inputs.Textbox(label="Model to Merge"),
+        gr.inputs.Textbox(label="Output Path")
     ],
     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'."
+    title="Model Merger",
+    description="Merge two models using the Super Mario merge method."
 )
 
-if __name__ == "__main__":
-    interface.launch()
+# Launch the interface
+iface.launch()