[TLM] Add reasoning effort option (#350)

Co-authored-by: Jay Zhang <[email protected]>

cleanlab_studio/internal/constants.py

@@ -30,6 +30,7 @@
 }
 TLM_NUM_CANDIDATE_RESPONSES_RANGE: Tuple[int, int] = (1, 20)  # (min, max)
 TLM_NUM_CONSISTENCY_SAMPLES_RANGE: Tuple[int, int] = (0, 20)  # (min, max)
+TLM_REASONING_EFFORT_VALUES: Set[str] = {"none", "low", "medium", "high"}
 TLM_VALID_LOG_OPTIONS: Set[str] = {"perplexity", "explanation"}
 TLM_VALID_GET_TRUSTWORTHINESS_SCORE_KWARGS: Set[str] = {"perplexity"}
 TLM_VALID_KWARGS: Set[str] = {"constrain_outputs"}

cleanlab_studio/internal/tlm/validation.py

@@ -8,6 +8,7 @@
     _VALID_TLM_MODELS,
     TLM_NUM_CANDIDATE_RESPONSES_RANGE,
     TLM_NUM_CONSISTENCY_SAMPLES_RANGE,
+    TLM_REASONING_EFFORT_VALUES,
     TLM_VALID_GET_TRUSTWORTHINESS_SCORE_KWARGS,
     TLM_VALID_KWARGS,
     TLM_VALID_LOG_OPTIONS,
@@ -231,6 +232,12 @@ def validate_tlm_options(options: Any) -> None:
                     f"Invalid type {type(val)}, use_self_reflection must be a boolean"
                 )
 
+        elif option == "reasoning_effort":
+            if val not in TLM_REASONING_EFFORT_VALUES:
+                raise ValidationError(
+                    f"Invalid value {val}, reasoning_effort must be one of {TLM_REASONING_EFFORT_VALUES}"
+                )
+
         elif option == "log":
             if not isinstance(val, list):
                 raise ValidationError(f"Invalid type {type(val)}, log must be a list of strings.")

cleanlab_studio/studio/trustworthy_language_model.py

@@ -822,6 +822,10 @@ class TLMOptions(TypedDict):
         and helping catch answers that are obviously incorrect/bad for a prompt asking for a well-defined answer that LLMs should be able to handle.
         Setting this to False disables the use of self-reflection and may produce worse TLM trustworthiness scores, but will reduce costs/runtimes.
 
+        reasoning_effort (str, default = "high"): Controls how much the LLM reasons when considering alternative possible responses and double-checking responses.
+        Higher efforts here produce better TLM trustworthiness scores, but at higher costs/runtimes, reduce this value to get faster results.
+        Supported efforts include "none", "low", "medium", "high".
+
         log (List[str], default = []): optionally specify additional logs or metadata to return.
         For instance, include "explanation" here to get explanations of why a response is scored with low trustworthiness.
 
@@ -837,6 +841,7 @@ class TLMOptions(TypedDict):
     num_candidate_responses: NotRequired[int]
     num_consistency_samples: NotRequired[int]
     use_self_reflection: NotRequired[bool]
+    reasoning_effort: NotRequired[str]
     log: NotRequired[List[str]]
     custom_eval_criteria: NotRequired[List[Dict[str, Any]]]
 

tests/tlm/conftest.py

@@ -12,6 +12,7 @@
     _TLM_MAX_TOKEN_RANGE,
     _VALID_TLM_MODELS,
     _VALID_TLM_QUALITY_PRESETS,
+    TLM_REASONING_EFFORT_VALUES,
 )
 from cleanlab_studio.internal.tlm.concurrency import TlmRateHandler
 from cleanlab_studio.studio.trustworthy_language_model import TLM
@@ -85,6 +86,7 @@ def _get_options_dictionary(model: Optional[str]) -> dict:
     add_num_candidate_responses = np.random.choice([True, False])
     add_num_consistency_samples = np.random.choice([True, False])
     add_use_self_reflection = np.random.choice([True, False])
+    add_reasoning_effort = np.random.choice([True, False])
     add_log_explanation = np.random.choice([True, False])
     add_log_perplexity_score = np.random.choice([True, False])
 
@@ -99,6 +101,8 @@ def _get_options_dictionary(model: Optional[str]) -> dict:
         options["num_candidate_responses"] = int(np.random.randint(1, 5))
     if add_num_consistency_samples:
         options["num_consistency_samples"] = int(np.random.randint(0, 10))
+    if add_reasoning_effort:
+        options["reasoning_effort"] = random.choice(list(TLM_REASONING_EFFORT_VALUES))
 
     if add_log_explanation or add_log_perplexity_score:
         options["log"] = [