From e087a4a72b9516f73b83d03662f70af9fb6471eb Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Tue, 10 Mar 2026 16:47:50 +0000
Subject: [PATCH] fix: improve Pass 2 prompt to correctly split INPUT/OUTPUT
 columns

The LLM was putting all columns in user_template (e.g. summarization
dataset had both document AND summary as user input). Fixed by:

- Reframed system message: explicitly states user=INPUT, assistant=OUTPUT
- Added 4 concrete correct examples (summarization, NLI, translation, QA)
  showing exactly how to split columns
- Added "NEVER put the output/target column in the user template" rule
- Added sanity check: if assistant_template has no column placeholders,
  reject the result and fall back to simple classification
---
 studio/backend/utils/datasets/llm_assist.py | 73 ++++++++++++++-------
 1 file changed, 49 insertions(+), 24 deletions(-)

diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py
index 4d6b6873f7..9963cf6bdf 100644
--- a/studio/backend/utils/datasets/llm_assist.py
+++ b/studio/backend/utils/datasets/llm_assist.py
@@ -545,15 +545,16 @@ def _run_multi_pass_advisor(
             {
                 "role": "system",
                 "content": (
-                    "You are a dataset conversion specialist for LLM fine-tuning. "
-                    "You design templates to convert non-conversational datasets into "
-                    "user/assistant conversation format. Respond with ONLY valid JSON."
+                    "You convert non-conversational datasets into user/assistant pairs "
+                    "for LLM fine-tuning. The user message is the INPUT (what the model "
+                    "receives). The assistant message is the OUTPUT (what the model should "
+                    "generate). You MUST have both. Respond with ONLY valid JSON."
                 ),
             },
             {
                 "role": "user",
                 "content": textwrap.dedent(f"""\
-                    This dataset was classified as:
+                    Dataset classification:
                     {json.dumps(pass1, indent=2)}
 
                     COLUMNS: {columns}
@@ -561,30 +562,41 @@ def _run_multi_pass_advisor(
                     SAMPLE DATA:
                     {samples_text}
 
-                    Design user and assistant message templates for this dataset.
+                    Split the columns into INPUT (user_template) and OUTPUT (assistant_template).
+                    Templates use ONLY {{column_name}} placeholders, NEVER actual data values.
+
+                    EXAMPLES of correct templates for different dataset types:
+                    - Summarization (columns: document, summary):
+                      user_template: "{{document}}"
+                      assistant_template: "{{summary}}"
+                      column_roles: {{"document": "user", "summary": "assistant"}}
+                    - NLI (columns: premise, hypothesis, label):
+                      user_template: "Premise: {{premise}}\\nHypothesis: {{hypothesis}}"
+                      assistant_template: "{{label_name}}"
+                      column_roles: {{"premise": "user", "hypothesis": "user", "label": "assistant"}}
+                    - Translation (columns: en, fr):
+                      user_template: "{{en}}"
+                      assistant_template: "{{fr}}"
+                      column_roles: {{"en": "user", "fr": "assistant"}}
+                    - QA (columns: question, context, answer):
+                      user_template: "Context: {{context}}\\nQuestion: {{question}}"
+                      assistant_template: "{{answer}}"
+                      column_roles: {{"context": "user", "question": "user", "answer": "assistant"}}
 
                     RULES:
-                    - Use {{column_name}} placeholders in templates to reference column values.
-                    - If a column has integer labels, provide a COMPLETE label_mapping for ALL values.
-                      Look at the actual sample data to determine what each integer means.
-                    - When a label_mapping exists for a column, use {{column_name_name}} in the
-                      assistant template to get the mapped string (not the raw integer).
-                    - The user template should include ALL relevant input columns.
-                    - The assistant template should produce the expected model output.
-                    - column_roles: mark columns used in the user template as "user",
-                      columns used in the assistant template as "assistant".
+                    - There MUST be at least one column as "user" AND at least one as "assistant".
+                    - NEVER put the output/target column in the user template.
+                    - If a column has integer labels, provide a label_mapping for ALL integer values.
+                    - When label_mapping exists for a column, use {{column_name}} in assistant_template
+                      (the mapping is applied automatically).
 
-                    Respond with a JSON object:
+                    Respond with JSON:
                     {{
-                        "user_template": "<template for user message using {{column}} placeholders>",
-                        "assistant_template": "<template for assistant response using {{column_name}} placeholders>",
-                        "column_roles": {{
-                            "<column_name>": "user or assistant (based on which template uses it)"
-                        }},
-                        "label_mapping": {{
-                            "<column_name>": {{"0": "<human-readable label>", "1": "<label>", ...}}
-                        }},
-                        "notes": "<any important notes about this conversion>"
+                        "user_template": "<INPUT template with {{column}} placeholders>",
+                        "assistant_template": "<OUTPUT template with {{column}} placeholders>",
+                        "column_roles": {{"<col>": "user or assistant"}},
+                        "label_mapping": {{"<col>": {{"0": "<label>", "1": "<label>"}}}},
+                        "notes": "<brief note>"
                     }}"""),
             },
         ]
@@ -601,6 +613,19 @@ def _run_multi_pass_advisor(
         asst_tpl = pass2.get("assistant_template", "")
         label_map = pass2.get("label_mapping", {})
 
+        # Sanity check: assistant_template must reference at least one column
+        has_asst_col = any(
+            f"{{{col}}}" in asst_tpl or f"{{{col}_name}}" in asst_tpl
+            for col in columns
+        )
+        if not has_asst_col and asst_tpl:
+            # LLM put literal text instead of a placeholder — reject
+            print(
+                f"🤖 Pass 2 sanity fail: assistant_template has no column placeholders: {asst_tpl!r}",
+                flush=True,
+            )
+            return None  # triggers fallback to simple classification
+
         # ── Pass 3: System prompt (non-conversational datasets only) ──
         sys_prompt = ""
         dtype = pass1.get("dataset_type", "unknown")