Add more UI components to trigger and to show the intermediate steps.

demo/sample_generator.ipynb

@@ -109,7 +109,7 @@
     "\"\"\")\n",
     "]\n",
     "\n",
-    "EXAMPLES_PROMPT = [\n",
+    "EXAMPLES_DIRECTLY_PROMPT = [\n",
     "    (\"system\", \"\"\"Given the task type description, and input/output example(s), generate {generating_batch_size}\n",
     "new input/output examples for this task type.\n",
     "\n",
@@ -150,7 +150,7 @@
     "        self.input_analysis_prompt = ChatPromptTemplate.from_messages(INPUT_ANALYSIS_PROMPT)\n",
     "        self.briefs_prompt = ChatPromptTemplate.from_messages(BRIEFS_PROMPT)\n",
     "        self.examples_from_briefs_prompt = ChatPromptTemplate.from_messages(EXAMPLES_FROM_BRIEFS_PROMPT)\n",
-    "        self.examples_prompt = ChatPromptTemplate.from_messages(EXAMPLES_PROMPT)\n",
+    "        self.examples_directly_prompt = ChatPromptTemplate.from_messages(EXAMPLES_DIRECTLY_PROMPT)\n",
     "\n",
     "        json_model = model.bind(response_format={\"type\": \"json_object\"})\n",
     "\n",
@@ -161,27 +161,34 @@
     "        self.input_analysis_chain = self.input_analysis_prompt | model | output_parser\n",
     "        self.briefs_chain = self.briefs_prompt | model | output_parser\n",
     "        self.examples_from_briefs_chain = self.examples_from_briefs_prompt | json_model | json_parse\n",
-    "        self.examples_chain = self.examples_prompt | json_model | json_parse\n",
+    "        self.examples_directly_chain = self.examples_directly_prompt | json_model | json_parse\n",
+    "\n",
+    "        # New sub-chain for loading and validating input\n",
+    "        self.input_loader = RunnableLambda(self.load_and_validate_input)\n",
     "\n",
     "        self.chain = (\n",
-    "            RunnablePassthrough.assign(raw_example = lambda x: json.dumps(x[\"example\"], ensure_ascii=False))\n",
+    "            self.input_loader\n",
+    "            | RunnablePassthrough.assign(raw_example = lambda x: json.dumps(x[\"example\"], ensure_ascii=False))\n",
     "            | RunnablePassthrough.assign(description = self.description_chain)\n",
     "            | {\n",
     "                \"description\": lambda x: x[\"description\"],\n",
     "                \"examples_from_briefs\": RunnablePassthrough.assign(input_analysis = self.input_analysis_chain)\n",
     "                    | RunnablePassthrough.assign(new_example_briefs = self.briefs_chain) \n",
     "                    | RunnablePassthrough.assign(examples = self.examples_from_briefs_chain | (lambda x: x[\"examples\"])),\n",
-    "                \"examples\": self.examples_chain\n",
+    "                \"examples_directly\": self.examples_directly_chain\n",
     "            }\n",
     "            | RunnablePassthrough.assign(\n",
     "                additional_examples=lambda x: (\n",
     "                    list(x[\"examples_from_briefs\"][\"examples\"])\n",
-    "                    + list(x[\"examples\"][\"examples\"])\n",
+    "                    + list(x[\"examples_directly\"][\"examples\"])\n",
     "                )\n",
     "            )\n",
     "        )\n",
     "\n",
-    "    def process(self, input_str, generating_batch_size=3):\n",
+    "    def load_and_validate_input(self, input_dict):\n",
+    "        input_str = input_dict[\"input_str\"]\n",
+    "        generating_batch_size = input_dict[\"generating_batch_size\"]\n",
+    "\n",
     "        try:\n",
     "            try:\n",
     "                example_dict = json.loads(input_str)\n",
@@ -202,15 +209,26 @@
     "            # Move the original content to a key named 'example'\n",
     "            input_dict = {\"example\": example_dict, \"generating_batch_size\": generating_batch_size}\n",
     "\n",
-    "            # Invoke the chain with the parsed input dictionary\n",
-    "            result = self.chain.invoke(input_dict)\n",
-    "            return result\n",
+    "            return input_dict\n",
     "\n",
     "        except Exception as e:\n",
     "            raise RuntimeError(f\"An error occurred during processing: {str(e)}\")\n",
     "\n",
-    "    def generate_description(self, input_str):\n",
-    "        return self.description_chain.invoke(input_str)\n",
+    "    def process(self, input_str, generating_batch_size=3):\n",
+    "        input_dict = {\"input_str\": input_str, \"generating_batch_size\": generating_batch_size}\n",
+    "        result = self.chain.invoke(input_dict)\n",
+    "        return result\n",
+    "\n",
+    "    def generate_description(self, input_str, generating_batch_size=3):\n",
+    "        chain = (\n",
+    "            self.input_loader \n",
+    "            | RunnablePassthrough.assign(raw_example = lambda x: json.dumps(x[\"example\"], ensure_ascii=False))\n",
+    "            | self.description_chain\n",
+    "        )\n",
+    "        return chain.invoke({\n",
+    "            \"input_str\": input_str,\n",
+    "            \"generating_batch_size\": generating_batch_size\n",
+    "        })\n",
     "\n",
     "    def analyze_input(self, description):\n",
     "        return self.input_analysis_chain.invoke(description)\n",
@@ -222,16 +240,25 @@
     "            \"generating_batch_size\": generating_batch_size\n",
     "        })\n",
     "\n",
-    "    def generate_examples_from_briefs(self, description, new_example_briefs, raw_example, generating_batch_size):\n",
-    "        return self.examples_from_briefs_chain.invoke({\n",
+    "    def generate_examples_from_briefs(self, description, new_example_briefs, input_str, generating_batch_size=3):\n",
+    "        chain = (\n",
+    "            self.input_loader\n",
+    "            | RunnablePassthrough.assign(\n",
+    "                raw_example = lambda x: json.dumps(x[\"example\"], ensure_ascii=False),\n",
+    "                description = lambda x: description,\n",
+    "                new_example_briefs = lambda x: new_example_briefs\n",
+    "            )\n",
+    "            | self.examples_from_briefs_chain\n",
+    "        )\n",
+    "        return chain.invoke({\n",
     "            \"description\": description,\n",
     "            \"new_example_briefs\": new_example_briefs,\n",
-    "            \"raw_example\": raw_example,\n",
+    "            \"input_str\": input_str,\n",
     "            \"generating_batch_size\": generating_batch_size\n",
     "        })\n",
     "\n",
-    "    def generate_examples(self, description, raw_example, generating_batch_size):\n",
-    "        return self.examples_chain.invoke({\n",
+    "    def generate_examples_directly(self, description, raw_example, generating_batch_size):\n",
+    "        return self.examples_directly_chain.invoke({\n",
     "            \"description\": description,\n",
     "            \"raw_example\": raw_example,\n",
     "            \"generating_batch_size\": generating_batch_size\n",
@@ -252,9 +279,12 @@
     "        generator = TaskDescriptionGenerator(model)\n",
     "        result = generator.process(input_json, generating_batch_size)\n",
     "        description = result[\"description\"]\n",
+    "        examples_directly = [[example[\"input\"], example[\"output\"]] for example in result[\"examples_directly\"][\"examples\"]]\n",
     "        input_analysis = result[\"examples_from_briefs\"][\"input_analysis\"]\n",
+    "        new_example_briefs = result[\"examples_from_briefs\"][\"new_example_briefs\"]\n",
+    "        examples_from_briefs = [[example[\"input\"], example[\"output\"]] for example in result[\"examples_from_briefs\"][\"examples\"]]\n",
     "        examples = [[example[\"input\"], example[\"output\"]] for example in result[\"additional_examples\"]]\n",
-    "        return description, input_analysis, examples\n",
+    "        return description, examples_directly, input_analysis, new_example_briefs, examples_from_briefs, examples\n",
     "    except Exception as e:\n",
     "        raise gr.Error(f\"An error occurred: {str(e)}\")\n",
     "    \n",
@@ -285,20 +315,22 @@
     "    except Exception as e:\n",
     "        raise gr.Error(f\"An error occurred: {str(e)}\")\n",
     "    \n",
-    "def generate_examples_from_briefs(description, new_example_briefs, raw_example, generating_batch_size, model_name, temperature):\n",
+    "def generate_examples_from_briefs(description, new_example_briefs, input_str, generating_batch_size, model_name, temperature):\n",
     "    try:\n",
     "        model = ChatOpenAI(model=model_name, temperature=temperature, max_retries=3)\n",
     "        generator = TaskDescriptionGenerator(model)\n",
-    "        examples = generator.generate_examples_from_briefs(description, new_example_briefs, raw_example, generating_batch_size)\n",
+    "        result = generator.generate_examples_from_briefs(description, new_example_briefs, input_str, generating_batch_size)\n",
+    "        examples = [[example[\"input\"], example[\"output\"]] for example in result[\"examples\"]]\n",
     "        return examples\n",
     "    except Exception as e:\n",
     "        raise gr.Error(f\"An error occurred: {str(e)}\")\n",
     "    \n",
-    "def generate_examples(description, raw_example, generating_batch_size, model_name, temperature):\n",
+    "def generate_examples_directly(description, raw_example, generating_batch_size, model_name, temperature):\n",
     "    try:\n",
     "        model = ChatOpenAI(model=model_name, temperature=temperature, max_retries=3)\n",
     "        generator = TaskDescriptionGenerator(model)\n",
-    "        examples = generator.generate_examples(description, raw_example, generating_batch_size)\n",
+    "        result = generator.generate_examples_directly(description, raw_example, generating_batch_size)\n",
+    "        examples = [[example[\"input\"], example[\"output\"]] for example in result[\"examples\"]]\n",
     "        return examples\n",
     "    except Exception as e:\n",
     "        raise gr.Error(f\"An error occurred: {str(e)}\")\n",
@@ -330,7 +362,10 @@
     "\n",
     "        with gr.Column(scale=1):  # Outputs column\n",
     "            description_output = gr.Textbox(label=\"Description\", lines=5, show_copy_button=True)\n",
-    "            analyze_input_button = gr.Button(\"Analyze Input\", variant=\"secondary\")\n",
+    "            with gr.Row():\n",
+    "                generate_examples_directly_button = gr.Button(\"Generate Examples Directly\", variant=\"secondary\")\n",
+    "                analyze_input_button = gr.Button(\"Analyze Input\", variant=\"secondary\")\n",
+    "            examples_directly_output = gr.DataFrame(label=\"Examples Directly\", headers=[\"Input\", \"Output\"], interactive=False)\n",
     "            input_analysis_output = gr.Textbox(label=\"Input Analysis\", lines=5, show_copy_button=True)\n",
     "            generate_briefs_button = gr.Button(\"Generate Briefs\", variant=\"secondary\")\n",
     "            example_briefs_output = gr.Textbox(label=\"Example Briefs\", lines=5, show_copy_button=True)\n",
@@ -346,13 +381,7 @@
     "    submit_button.click(\n",
     "        fn=process_json,\n",
     "        inputs=[input_json, model_name, generating_batch_size, temperature],\n",
-    "        outputs=[description_output, input_analysis_output, examples_output]\n",
-    "    )\n",
-    "\n",
-    "    analyze_input_button.click(\n",
-    "        fn=analyze_input,\n",
-    "        inputs=[description_output, model_name, temperature],\n",
-    "        outputs=[input_analysis_output]\n",
+    "        outputs=[description_output, examples_directly_output, input_analysis_output, example_briefs_output, examples_from_briefs_output, examples_output]\n",
     "    )\n",
     "\n",
     "    generate_description_button.click(\n",
@@ -361,6 +390,18 @@
     "        outputs=[description_output]\n",
     "    )\n",
     "\n",
+    "    generate_examples_directly_button.click(\n",
+    "        fn=generate_examples_directly,\n",
+    "        inputs=[description_output, input_json, generating_batch_size, model_name, temperature],\n",
+    "        outputs=[examples_directly_output]\n",
+    "    )\n",
+    "\n",
+    "    analyze_input_button.click(\n",
+    "        fn=analyze_input,\n",
+    "        inputs=[description_output, model_name, temperature],\n",
+    "        outputs=[input_analysis_output]\n",
+    "    )\n",
+    "\n",
     "    generate_briefs_button.click(\n",
     "        fn=generate_briefs,\n",
     "        inputs=[description_output, input_analysis_output, generating_batch_size, model_name, temperature],\n",
@@ -373,6 +414,18 @@
     "        outputs=[examples_from_briefs_output]\n",
     "    )\n",
     "\n",
+    "    examples_directly_output.select(\n",
+    "        fn=format_selected_example,\n",
+    "        inputs=[examples_directly_output],\n",
+    "        outputs=[new_example_json]\n",
+    "    )\n",
+    "\n",
+    "    examples_from_briefs_output.select(\n",
+    "        fn=format_selected_example,\n",
+    "        inputs=[examples_from_briefs_output],\n",
+    "        outputs=[new_example_json]\n",
+    "    )\n",
+    "\n",
     "    examples_output.select(\n",
     "        fn=format_selected_example,\n",
     "        inputs=[examples_output],\n",