mvp

app.py

@@ -1,30 +1,89 @@
 import gradio as gr
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+import torch
+import torch.distributed.run as distributed_run
+from git import Repo
+from huggingface_hub import HfApi
+
+
+# Clone the medusa repo locally
+Repo.clone_from("https://github.com/FasterDecoding/Medusa.git", "medusa")
+
 
 def create_medusa_heads(model_id: str):
-    return "", ""
+    parser = distributed_run.get_args_parser()
+    args = parser.parse_args([
+        "--nproc_per_node", "4",
+        "--training_script", "medusa/medusa/train/train.py",
+        "--training_script_args",
+        "--model_name_or_path", model_id,
+        "--data_path", "ShareGPT_Vicuna_unfiltered/ShareGPT_V4.3_unfiltered_cleaned_split.json",
+        "--bf16", "True",
+        "--output_dir", "medusa_heads",
+        "--num_train_epochs", "1",
+        "--per_device_train_batch_size", "8",
+        "--per_device_eval_batch_size", "8",
+        "--gradient_accumulation_steps", "4",
+        "--evaluation_strategy", "no",
+        "--save_strategy", "no",
+        "--learning_rate", "1e-3",
+        "--weight_decay", "0.0",
+        "--warmup_ratio", "0.1",
+        "--lr_scheduler_type", "cosine",
+        "--logging_steps", "1",
+        "--tf32", "True",
+        "--model_max_length", "2048",
+        "--lazy_preprocess", "True",
+        "--medusa_num_heads", "3",
+        "--medusa_num_layers", "1",
+    ])
+    distributed_run.run(args)
+
+    # Upload the medusa heads to the Hub
+    repo_id = f"medusa-{model_id}"
+    api = HfApi()
+    api.create_repo(
+        repo_id=repo_id,
+        exist_ok=True,
+    )
+    api.upload_folder(
+        folder_path="medusa_heads",
+        repo_id=repo_id,
+    )
+    return repo_id
 
 def run(model_id: str) -> str:
+    # Input validation
     if model_id == "":
         return """
         ### Invalid input 🐞
 
         Please fill a model_id.
         """
+
+    # Attempt to load the base model
     try:
-        commit_info, errors = create_medusa_heads(model_id=model_id)
-        print("[commit_info]", commit_info)
+        config = AutoConfig.from_pretrained(model_id)
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)
+        del config, tokenizer, model
+    except Exception as e:
+        return f"""
+        ### {model_id} can't be loaded with AutoClasses 🐞
 
-        string =  f"""
-        ### Success 🔥
+        {e}
+        """
+    print(f"{model_id} can be loaded, starting medusa heads creation.")
 
-        Yay! This model was successfully converted and a PR was open using your token, here:
+    # Run the medusa heads creation
+    try:
+        repo_id = create_medusa_heads(model_id=model_id)
+        print("Medusa heads uploaded to: ", repo_id)
+        return f"""
+        ### Success 🔥
 
-        [{commit_info.pr_url}]({commit_info.pr_url})
+        Yay! Medusa heads were successfully created and uploaded to, {repo_id}
         """
-        if errors:
-            string += "\nErrors during conversion:\n"
-            string += "\n".join(f"Error while converting {filename}: {e}, skipped conversion" for filename, e in errors)
-        return string
     except Exception as e:
         return f"""
         ### Error 😢😢😢
@@ -34,11 +93,11 @@ def run(model_id: str) -> str:
 
 
 DESCRIPTION = """
-The step to create [medusa](https://sites.google.com/view/medusa-llm) heads are the following:
+The steps to create [medusa](https://sites.google.com/view/medusa-llm) heads are the following:
 
-- Input a public model id from the Hub
-- Click "Submit"
-- That's it! You'll get feedback if it works or not, and if it worked, you'll get the URL of the new repo 🔥
+1. Input a public model id from the Hub
+2. Click "Submit"
+3. That's it! You'll get feedback if it works or not, and if it worked, you'll get the URL of the new repo 🔥
 """
 
 title="Create LLM medusa heads in a new repo 🐍"

requirements.txt

@@ -1 +1,3 @@
-medusa-llm[train]
+gitpython
+transformers
+torch