Forgot docker launcher.

This commit is contained in:
Nicolas Patry 2024-02-26 17:07:54 +00:00
parent bfec09ecc2
commit 915e5f088c
2 changed files with 5 additions and 5 deletions

View File

@ -306,6 +306,7 @@ def launcher(event_loop):
use_flash_attention: bool = True,
disable_grammar_support: bool = False,
dtype: Optional[str] = None,
revision: Optional[str] = None,
):
port = random.randint(8000, 10_000)
@ -321,6 +322,9 @@ def launcher(event_loop):
if dtype is not None:
args.append("--dtype")
args.append(dtype)
if revision is not None:
args.append("--revision")
args.append(revision)
if trust_remote_code:
args.append("--trust-remote-code")

View File

@ -154,12 +154,8 @@ def download_weights(
import json
medusa_head = hf_hub_download(
model_id, revision=revision, filename="medusa_lm_head.pt"
model_id, revision=revision, filename="medusa_lm_head.safetensors"
)
if auto_convert:
medusa_sf = Path(medusa_head[: -len(".pt")] + ".safetensors")
if not medusa_sf.exists():
utils.convert_files([Path(medusa_head)], [medusa_sf], [])
medusa_config = hf_hub_download(
model_id, revision=revision, filename="config.json"
)