From 5db251b31c1b2c997bdea78cb50672ca310e2347 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 16:21:54 +0000 Subject: [PATCH] fix: include label mapping in Pass 3 system prompt generation Pass 3 now sees the label mapping from Pass 2 (e.g. "0 = does not follow, 1 = follows, 2 = entailed") so the generated system prompt can explain what each label value means. Also bumped to 2-4 sentences to give room for the label descriptions. --- studio/backend/utils/datasets/llm_assist.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index b5f8ba6b35..4d6b6873f7 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -609,6 +609,15 @@ def _run_multi_pass_advisor( if not is_conv: print("🤖 Pass 3: Generating system prompt...", flush=True) t3 = time.monotonic() + + # Format label mapping info for the prompt + label_info = "" + if label_map: + for col, mapping in label_map.items(): + if isinstance(mapping, dict) and mapping: + pairs = ", ".join(f"{k} = {v}" for k, v in mapping.items()) + label_info += f"\nLabel mapping for '{col}': {pairs}" + messages3 = [ { "role": "system", @@ -629,13 +638,15 @@ def _run_multi_pass_advisor( The training examples will look like: User: {user_tpl} Assistant: {asst_tpl} + {label_info} Write a system prompt that clearly describes the task the model should perform. The system prompt should: - Explain what kind of input the model will receive - Explain what output is expected + - If there are categorical labels, explain what each label means - Be specific to this dataset's task (not generic) - - Be 1-3 sentences + - Be 2-4 sentences Respond with a JSON object: {{