joaogante HF staff commited on
Commit
b6a509f
β€’
1 Parent(s): 6c981de
Files changed (2) hide show
  1. app.py +74 -15
  2. requirements.txt +3 -1
app.py CHANGED
@@ -1,30 +1,89 @@
1
  import gradio as gr
 
 
 
 
 
 
 
 
 
 
2
 
3
  def create_medusa_heads(model_id: str):
4
- return "", ""
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5
 
6
  def run(model_id: str) -> str:
 
7
  if model_id == "":
8
  return """
9
  ### Invalid input 🐞
10
 
11
  Please fill a model_id.
12
  """
 
 
13
  try:
14
- commit_info, errors = create_medusa_heads(model_id=model_id)
15
- print("[commit_info]", commit_info)
 
 
 
 
 
16
 
17
- string = f"""
18
- ### Success πŸ”₯
 
19
 
20
- Yay! This model was successfully converted and a PR was open using your token, here:
 
 
 
 
 
21
 
22
- [{commit_info.pr_url}]({commit_info.pr_url})
23
  """
24
- if errors:
25
- string += "\nErrors during conversion:\n"
26
- string += "\n".join(f"Error while converting {filename}: {e}, skipped conversion" for filename, e in errors)
27
- return string
28
  except Exception as e:
29
  return f"""
30
  ### Error 😒😒😒
@@ -34,11 +93,11 @@ def run(model_id: str) -> str:
34
 
35
 
36
  DESCRIPTION = """
37
- The step to create [medusa](https://sites.google.com/view/medusa-llm) heads are the following:
38
 
39
- - Input a public model id from the Hub
40
- - Click "Submit"
41
- - 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 πŸ”₯
42
  """
43
 
44
  title="Create LLM medusa heads in a new repo 🐍"
 
1
  import gradio as gr
2
+ from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
3
+ import torch
4
+ import torch.distributed.run as distributed_run
5
+ from git import Repo
6
+ from huggingface_hub import HfApi
7
+
8
+
9
+ # Clone the medusa repo locally
10
+ Repo.clone_from("https://github.com/FasterDecoding/Medusa.git", "medusa")
11
+
12
 
13
  def create_medusa_heads(model_id: str):
14
+ parser = distributed_run.get_args_parser()
15
+ args = parser.parse_args([
16
+ "--nproc_per_node", "4",
17
+ "--training_script", "medusa/medusa/train/train.py",
18
+ "--training_script_args",
19
+ "--model_name_or_path", model_id,
20
+ "--data_path", "ShareGPT_Vicuna_unfiltered/ShareGPT_V4.3_unfiltered_cleaned_split.json",
21
+ "--bf16", "True",
22
+ "--output_dir", "medusa_heads",
23
+ "--num_train_epochs", "1",
24
+ "--per_device_train_batch_size", "8",
25
+ "--per_device_eval_batch_size", "8",
26
+ "--gradient_accumulation_steps", "4",
27
+ "--evaluation_strategy", "no",
28
+ "--save_strategy", "no",
29
+ "--learning_rate", "1e-3",
30
+ "--weight_decay", "0.0",
31
+ "--warmup_ratio", "0.1",
32
+ "--lr_scheduler_type", "cosine",
33
+ "--logging_steps", "1",
34
+ "--tf32", "True",
35
+ "--model_max_length", "2048",
36
+ "--lazy_preprocess", "True",
37
+ "--medusa_num_heads", "3",
38
+ "--medusa_num_layers", "1",
39
+ ])
40
+ distributed_run.run(args)
41
+
42
+ # Upload the medusa heads to the Hub
43
+ repo_id = f"medusa-{model_id}"
44
+ api = HfApi()
45
+ api.create_repo(
46
+ repo_id=repo_id,
47
+ exist_ok=True,
48
+ )
49
+ api.upload_folder(
50
+ folder_path="medusa_heads",
51
+ repo_id=repo_id,
52
+ )
53
+ return repo_id
54
 
55
  def run(model_id: str) -> str:
56
+ # Input validation
57
  if model_id == "":
58
  return """
59
  ### Invalid input 🐞
60
 
61
  Please fill a model_id.
62
  """
63
+
64
+ # Attempt to load the base model
65
  try:
66
+ config = AutoConfig.from_pretrained(model_id)
67
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
68
+ model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)
69
+ del config, tokenizer, model
70
+ except Exception as e:
71
+ return f"""
72
+ ### {model_id} can't be loaded with AutoClasses 🐞
73
 
74
+ {e}
75
+ """
76
+ print(f"{model_id} can be loaded, starting medusa heads creation.")
77
 
78
+ # Run the medusa heads creation
79
+ try:
80
+ repo_id = create_medusa_heads(model_id=model_id)
81
+ print("Medusa heads uploaded to: ", repo_id)
82
+ return f"""
83
+ ### Success πŸ”₯
84
 
85
+ Yay! Medusa heads were successfully created and uploaded to, {repo_id}
86
  """
 
 
 
 
87
  except Exception as e:
88
  return f"""
89
  ### Error 😒😒😒
 
93
 
94
 
95
  DESCRIPTION = """
96
+ The steps to create [medusa](https://sites.google.com/view/medusa-llm) heads are the following:
97
 
98
+ 1. Input a public model id from the Hub
99
+ 2. Click "Submit"
100
+ 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 πŸ”₯
101
  """
102
 
103
  title="Create LLM medusa heads in a new repo 🐍"
requirements.txt CHANGED
@@ -1 +1,3 @@
1
- medusa-llm[train]
 
 
 
1
+ gitpython
2
+ transformers
3
+ torch