Adding trust_remote_code to the orchestrator and worker
This commit is contained in:
parent
86e94b5844
commit
662cb1c440
2 changed files with 4 additions and 0 deletions
|
|
@ -211,6 +211,7 @@ class ExportOrchestrator:
|
|||
checkpoint_path: str,
|
||||
max_seq_length: int = 2048,
|
||||
load_in_4bit: bool = True,
|
||||
trust_remote_code: bool = False,
|
||||
) -> Tuple[bool, str]:
|
||||
"""Load a checkpoint for export.
|
||||
|
||||
|
|
@ -225,6 +226,7 @@ class ExportOrchestrator:
|
|||
"checkpoint_path": checkpoint_path,
|
||||
"max_seq_length": max_seq_length,
|
||||
"load_in_4bit": load_in_4bit,
|
||||
"trust_remote_code": trust_remote_code,
|
||||
}
|
||||
|
||||
# Always kill existing subprocess and spawn fresh.
|
||||
|
|
|
|||
|
|
@ -84,6 +84,7 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
|
|||
checkpoint_path = cmd["checkpoint_path"]
|
||||
max_seq_length = cmd.get("max_seq_length", 2048)
|
||||
load_in_4bit = cmd.get("load_in_4bit", True)
|
||||
trust_remote_code = cmd.get("trust_remote_code", False)
|
||||
|
||||
try:
|
||||
_send_response(resp_queue, {
|
||||
|
|
@ -96,6 +97,7 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
|
|||
checkpoint_path=checkpoint_path,
|
||||
max_seq_length=max_seq_length,
|
||||
load_in_4bit=load_in_4bit,
|
||||
trust_remote_code=trust_remote_code,
|
||||
)
|
||||
|
||||
_send_response(resp_queue, {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue