From 3842b00736c3d849d8089bd0be77393be0059d00 Mon Sep 17 00:00:00 2001
From: pelikhan <jhalleux@microsoft.com>
Date: Fri, 17 Jan 2025 11:41:01 -0800
Subject: [PATCH] support disabling explanataions

---
 .../content/docs/reference/scripts/classify.md   | 15 ++++++++++++++-
 packages/cli/src/runtime.ts                      | 16 ++++++++++++----
 packages/sample/genaisrc/classify.genai.mjs      |  3 ++-
 3 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/docs/src/content/docs/reference/scripts/classify.md b/docs/src/content/docs/reference/scripts/classify.md
index 61c648eeee..fbe6aa0b34 100644
--- a/docs/src/content/docs/reference/scripts/classify.md
+++ b/docs/src/content/docs/reference/scripts/classify.md
@@ -52,7 +52,7 @@ Each label id should be a single word that encodes into a single token. This all
 A `other` label can be automatically added to the list
 of label to give an escape route for the LLM when it is not able to classify the text.
 
-```js
+```js "other: true"
 const res = await classify(
     "...",
     { ... },
@@ -60,6 +60,19 @@ const res = await classify(
 )
 ```
 
+## Explanations
+
+By default, the classification prompt is tuned to return a token (`maxToken: 1`) as the label.
+You can enable emitting a justification before returning the label.
+
+```js "explanation: true"
+const res = await classify(
+    "...",
+    { ... },
+    { explanation: true }
+)
+```
+
 ## Model and other options
 
 The `classify` function uses the `classify` [model alias](/genaiscript/reference/scripts/model-aliases) by default.
diff --git a/packages/cli/src/runtime.ts b/packages/cli/src/runtime.ts
index baa4626e08..492573f8d5 100644
--- a/packages/cli/src/runtime.ts
+++ b/packages/cli/src/runtime.ts
@@ -16,6 +16,13 @@ export type ClassifyOptions = {
      * Inject a 'other' label
      */
     other?: boolean
+    /**
+     * Explain answers before returning token
+     */
+    explanations?: boolean
+    /**
+     * Options runPrompt context
+     */
     ctx?: ChatGenerationContext
 } & Omit<PromptGeneratorOptions, "choices">
 
@@ -40,7 +47,7 @@ export async function classify<L extends Record<string, string>>(
     answer: string
     logprobs?: Record<keyof typeof labels | "other", Logprob>
 }> {
-    const { other, ...rest } = options || {}
+    const { other, explanations, ...rest } = options || {}
 
     const entries = Object.entries({
         ...labels,
@@ -71,8 +78,8 @@ You must classify the data as one of the following labels.
 ${entries.map(([id, descr]) => `- Label '${id}': ${descr}`).join("\n")}
 
 ## Output
-Provide a single sentence justification for your choice.
-and output the label as a single word on the last line. Do not emit "Label".
+${explanations ? "Provide a single short sentence justification for your choice." : ""}
+Output the label as a single word on the last line (do not emit "Label").
 
 `
             _.fence(
@@ -83,7 +90,7 @@ DATA:
 Why did the chicken cross the road? Because moo.
 
 Output:
-It's a classic joke but the ending does not relate to the start of the joke.
+${explanations ? "It's a classic joke but the ending does not relate to the start of the joke." : ""}
 no
 
 `,
@@ -98,6 +105,7 @@ no
             label: `classify ${choices.join(", ")}`,
             logprobs: true,
             topLogprobs: Math.min(3, choices.length),
+            maxTokens: explanations ? 100 : 1,
             system: [
                 "system.output_plaintext",
                 "system.safety_jailbreak",
diff --git a/packages/sample/genaisrc/classify.genai.mjs b/packages/sample/genaisrc/classify.genai.mjs
index 79eb373912..1ff3da7e99 100644
--- a/packages/sample/genaisrc/classify.genai.mjs
+++ b/packages/sample/genaisrc/classify.genai.mjs
@@ -17,7 +17,8 @@ const joke = await classify(
     {
         yes: "funny",
         no: "not funny",
-    }
+    },
+    { explanations: true }
 )
 
 console.log(joke)