Parallel run.

demo/sample_generator.ipynb

@@ -63,6 +63,7 @@
     "from langchain.prompts import ChatPromptTemplate\n",
     "from langchain.chat_models import ChatOpenAI\n",
     "from langchain.schema.output_parser import StrOutputParser\n",
+    "from langchain_core.runnables import RunnablePassthrough, RunnableParallel\n",
     "\n",
     "class TaskDescriptionGenerator:\n",
     "    def __init__(self, model_temperature=0.7):\n",
@@ -71,56 +72,56 @@
     "        \n",
     "        self.description_prompt = ChatPromptTemplate.from_template(\n",
     "            \"\"\"Given the following JSON example for a task type:\n",
-    "            {raw_example}\n",
+    "{raw_example}\n",
     "\n",
-    "            Provide a concise description of the task type, including the format and style of the output.\n",
+    "Provide a concise description of the task type, including the format and style of the output.\n",
     "\n",
-    "            Format your response as follows:\n",
-    "            Task Description: [Your description here]\n",
-    "            \"\"\"\n",
+    "Format your response as follows:\n",
+    "Task Description: [Your description here]\n",
+    "\"\"\"\n",
     "        )\n",
     "\n",
     "        self.briefs_prompt = ChatPromptTemplate.from_template(\n",
     "            \"\"\"Given the following task type description:\n",
-    "            Task Description: {description}\n",
+    "Task Description: {description}\n",
     "\n",
-    "            Generate descriptions for 3 new examples with detailed attributes\n",
-    "            based on this task type.\n",
+    "Generate descriptions for 3 new examples with detailed attributes\n",
+    "based on this task type.\n",
     "\n",
-    "            Format your response as a valid YAML object with a single key\n",
-    "            'brief_descriptions' containing a YAML array of 3 objects, each\n",
-    "            with a 'description' field.\n",
-    "            \"\"\"\n",
+    "Format your response as a valid YAML object with a single key\n",
+    "'brief_descriptions' containing a YAML array of 3 objects, each\n",
+    "with a 'description' field.\n",
+    "\"\"\"\n",
     "        )\n",
     "\n",
     "        self.examples_from_briefs_prompt = ChatPromptTemplate.from_template(\n",
     "            \"\"\"Given the following task type description, brief description for new examples, \n",
-    "            and JSON example:\n",
+    "and JSON example:\n",
     "\n",
-    "            Task Description: {description}\n",
-    "            Brief Description: {brief_description}\n",
+    "Task Description: {description}\n",
+    "Brief Description: {brief_description}\n",
     "\n",
-    "            Example:\n",
-    "            {raw_example}\n",
+    "Example:\n",
+    "{raw_example}\n",
     "\n",
-    "            Generate 3 more input/output examples for this task type, based on the brief\n",
-    "            description, in the same JSON format.\n",
+    "Generate 3 more input/output examples for this task type, based on the brief\n",
+    "description, in the same JSON format.\n",
     "\n",
-    "            Format your response as a valid JSON object with a single key 'examples' \n",
-    "            containing a JSON array of 3 objects, each with 'input' and 'output' fields.\n",
-    "            \"\"\"\n",
+    "Format your response as a valid JSON object with a single key 'examples' \n",
+    "containing a JSON array of 3 objects, each with 'input' and 'output' fields.\n",
+    "\"\"\"\n",
     "        )\n",
     "\n",
     "        self.examples_prompt = ChatPromptTemplate.from_template(\n",
     "            \"\"\"Given the following task type description, and input/output example:\n",
-    "            Task Description: {description}\n",
-    "            Example: {example}\n",
+    "Task Description: {description}\n",
+    "Example: {raw_example}\n",
     "\n",
-    "            Generate 3 new input/output examples for this task type.\n",
+    "Generate 3 new input/output examples for this task type.\n",
     "\n",
-    "            Format your response as a valid JSON object with a single key 'examples' \n",
-    "            containing a JSON array of 3 objects, each with 'input' and 'output' fields.\n",
-    "            \"\"\"\n",
+    "Format your response as a valid JSON object with a single key 'examples' \n",
+    "containing a JSON array of 3 objects, each with 'input' and 'output' fields.\n",
+    "\"\"\"\n",
     "        )\n",
     "\n",
     "        self.description_chain = self.description_prompt | self.model | self.output_parser\n",
@@ -178,35 +179,45 @@
     "\n",
     "    def process(self, input_str):\n",
     "        try:\n",
-    "            # Try to parse as JSON\n",
-    "            try:\n",
-    "                data = json.loads(input_str)\n",
-    "            except json.JSONDecodeError:\n",
-    "                # If JSON parsing fails, try to parse as YAML\n",
-    "                try:\n",
-    "                    data = yaml.safe_load(input_str)\n",
-    "                except yaml.YAMLError as e:\n",
-    "                    raise ValueError(\"Invalid input string. Error: {}\".format(str(e)))\n",
-    "            \n",
-    "            if not isinstance(data, dict) or 'input' not in data or 'output' not in data:\n",
-    "                raise ValueError(\"Invalid input format. Expected an object with 'input' and 'output' fields.\")\n",
-    "            \n",
-    "            description = self.generate_description(json.dumps(data, ensure_ascii=False))\n",
-    "            brief_description = self.generate_briefs(description)\n",
-    "            new_examples_from_briefs = self.generate_examples_from_briefs(description, brief_description, json.dumps(data, ensure_ascii=False))\n",
-    "            new_examples = self.generate_examples(description, json.dumps(data, ensure_ascii=False))\n",
-    "    \n",
-    "            output = {\n",
-    "                \"task_description\": description,\n",
-    "                \"brief_description\": brief_description,\n",
-    "                \"additional_examples\": list(new_examples_from_briefs.values()) + list(new_examples.values())\n",
-    "            }\n",
-    "            \n",
-    "            return json.dumps(output, indent=2, ensure_ascii=False)\n",
-    "        \n",
+    "            # Parse input\n",
+    "            parse_input = (\n",
+    "                RunnablePassthrough()\n",
+    "                | (lambda x: json.loads(x))\n",
+    "                | (lambda x: x if isinstance(x, dict) and 'input' in x and 'output' in x \n",
+    "                   else (_ for _ in ()).throw(ValueError(\"Invalid input format. Expected an object with 'input' and 'output' fields.\")))\n",
+    "            )\n",
+    "\n",
+    "            # Create the chain\n",
+    "            chain = (\n",
+    "                parse_input\n",
+    "                | (lambda x: json.dumps(x, ensure_ascii=False))\n",
+    "                | {\n",
+    "                    \"description\": self.description_chain,\n",
+    "                    \"raw_example\": lambda x: x\n",
+    "                }\n",
+    "                | {\n",
+    "                    \"examples_from_briefs\": {\n",
+    "                        \"description\": lambda x: x[\"description\"],\n",
+    "                        \"brief_description\": lambda x: self.briefs_chain.invoke({\"description\": x[\"description\"]}),\n",
+    "                        \"raw_example\": lambda x: x[\"raw_example\"]\n",
+    "                    } | self.examples_from_briefs_chain,\n",
+    "                    \"examples\": self.examples_chain\n",
+    "                }\n",
+    "                | {\n",
+    "                    \"additional_examples\": lambda x: (\n",
+    "                        list(json.loads(x[\"examples_from_briefs\"])[\"examples\"])\n",
+    "                        + list(json.loads(x[\"examples\"])[\"examples\"])\n",
+    "                    )\n",
+    "                }\n",
+    "                | (lambda x: json.dumps(x, indent=2, ensure_ascii=False))\n",
+    "            )\n",
+    "\n",
+    "            # Invoke the chain\n",
+    "            result = chain.invoke(input_str)\n",
+    "            return result\n",
+    "\n",
     "        except Exception as e:\n",
     "            raise RuntimeError(f\"An error occurred during processing: {str(e)}\")\n",
-    "    \n",
     "\n",
     "import gradio as gr\n",
     "\n",