fdsp config

app.py

@@ -1,3 +1,4 @@
+import json
 import os
 import multiprocessing as mp
 
@@ -54,7 +55,7 @@ DEFAULT_TRAINING_ARGS = \
 --auto_find_batch_size True"""
 
 
-def train_medusa_heads(model_id: str, training_args: str):
+def train_medusa_heads(model_id: str, training_args: str, fdsp_config: str):
     all_training_args = FIXED_TRAINING_ARGS.format(
         model_id=model_id,
         output_dir=OUTPUT_DIR,
@@ -67,13 +68,16 @@ def train_medusa_heads(model_id: str, training_args: str):
     for arg in all_training_args.split("\n"):
         all_training_arg_list += arg.split(" ")
 
-    parser = distributed_run.get_args_parser()
+    if fdsp_config != "":
+        all_training_arg_list += ["--fdsp_config", json.loads(fdsp_config)]
     print(all_training_arg_list)
+
+    parser = distributed_run.get_args_parser()
     args = parser.parse_args(all_training_arg_list)
     distributed_run.run(args)
 
 
-def run(model_id: str, training_args: str) -> str:
+def run(model_id: str, training_args: str, fdsp_config: str) -> str:
     print(f"\n\n\nNEW RUN: {model_id}")
     api = HfApi()
     model_name = model_id.split("/")[-1]
@@ -110,10 +114,10 @@ def run(model_id: str, training_args: str) -> str:
 
     # Run the medusa heads creation
     try:
-        proc = mp.Process(target=train_medusa_heads, args=(model_id, training_args))
+        proc = mp.Process(target=train_medusa_heads, args=(model_id, training_args, fdsp_config))
         proc.start()
         proc.join()
-        print("Medusa heads training success ✅")
+        print("Medusa heads training process completed (it might have crashed!)")
     except Exception as e:
         print("Error ❌\n", e)
         return f"""
@@ -178,8 +182,14 @@ with gr.Blocks(title=title) as demo:
     with gr.Row() as r:
         with gr.Column() as c:
             model_id = gr.Text(max_lines=1, label="model_id")
-            with gr.Accordion("Training arguments", open=False):
+            with gr.Accordion("Training arguments (advanced)", open=False):
                 training_args = gr.Textbox(DEFAULT_TRAINING_ARGS, interactive=True, lines=14, label="training_args")
+                fdsp_config = gr.Textbox(
+                    placeholder="e.g. \'{\"value1\":\"key1\"}\'. leave empty if fdsp is not used. ",
+                    interactive=True,
+                    lines=1,
+                    label="fdsp_config"
+                )
             with gr.Row() as c:
                 clean = gr.ClearButton()
                 submit = gr.Button("Submit", variant="primary")
@@ -187,6 +197,6 @@ with gr.Blocks(title=title) as demo:
         with gr.Column() as d:
             status_box = gr.Markdown()
 
-    submit.click(run, inputs=[model_id, training_args], outputs=status_box, concurrency_limit=1)
+    submit.click(run, inputs=[model_id, training_args, fdsp_config], outputs=status_box, concurrency_limit=1)
 
 demo.queue(max_size=10).launch(show_api=True)