更新medusa权重转换脚本

examples/medusa/medusa_weight_converter.py

@@ -369,7 +369,7 @@ def main(args):
         os.makedirs(args.output_dir, exist_ok=True)
     save_model(medusa_model, os.path.join(args.output_dir, "model.safetensors"))
     
-    medusa_choices = ast.literal_eval(args.medusa_choices)
+    medusa_choices = ast.literal_eval(args.medusa_choices) if args.medusa_choices is not None else None
     to_save_config = CustomMedusaConfig(name_or_path=os.path.join(args.output_dir, "config.json"),
                                         hidden_size=args.hidden_size,
                                         num_heads=medusa_head_num,
@@ -403,7 +403,7 @@ if __name__ == "__main__":
     parser.add_argument(
         '--medusa_choices',
         type=str,
-        required=True,
+        default=None,
         help="Medusa choice to use, if not none, will use Medusa decoding."
         "   E.g.: [[0, 0, 0, 0], [0, 1, 0], [1, 0], [1, 1]] for 9 medusa tokens."
     )