diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index ec18d3533e..790ac4b000 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -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. diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index 9e7e72e9dd..82428e670b 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -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, {