File tree Expand file tree Collapse file tree 1 file changed +11
-0
lines changed
Expand file tree Collapse file tree 1 file changed +11
-0
lines changed Original file line number Diff line number Diff line change 55
66import mlflow
77from modelgauge .pipeline_runner import build_runner
8+ from modelgauge .sut import SUTOptions
89from modelgauge .sut_factory import SUT_FACTORY
910
1011from modelplane .runways .data import (
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
2736def 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 (
You can’t perform that action at this time.
0 commit comments