Skip to content

Commit 8667e56

Browse files
committed
Add DEFAULT_SUT_OPTIONS and pass sut_options to respond function.
1 parent f1fd7ed commit 8667e56

File tree

1 file changed

+11
-0
lines changed

1 file changed

+11
-0
lines changed

src/modelplane/runways/responder.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
import mlflow
77
from modelgauge.pipeline_runner import build_runner
8+
from modelgauge.sut import SUTOptions
89
from modelgauge.sut_factory import SUT_FACTORY
910

1011
from modelplane.runways.data import (
@@ -23,6 +24,14 @@
2324
setup_sut_credentials,
2425
)
2526

27+
# Taken from modelgauge.tests.security.BaseSecurityTest
28+
# TODO: Figure out a way to expose the options in the CLI.
29+
DEFAULT_SUT_OPTIONS = SUTOptions(
30+
max_tokens=3000,
31+
max_total_output_tokens=10000, # For reasoning SUTs.
32+
temperature=0.01,
33+
)
34+
2635

2736
def respond(
2837
sut_id: str,
@@ -34,6 +43,7 @@ def respond(
3443
num_workers: int = 1,
3544
prompt_uid_col=None,
3645
prompt_text_col=None,
46+
sut_options: SUTOptions = DEFAULT_SUT_OPTIONS,
3747
) -> RunArtifacts:
3848
secrets = setup_sut_credentials(sut_id)
3949
sut = SUT_FACTORY.make_instance(uid=sut_id, secrets=secrets)
@@ -60,6 +70,7 @@ def respond(
6070
suts={sut.uid: sut},
6171
prompt_uid_col=prompt_uid_col,
6272
prompt_text_col=prompt_text_col,
73+
sut_options=sut_options,
6374
)
6475

6576
pipeline_runner.run(

0 commit comments

Comments
 (0)