From f4ef67c5d3723e84c970fe221b6a49f9ada8962a Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Sat, 30 Mar 2024 16:29:27 -0700 Subject: [PATCH] Draft tnt llm --- examples/tnt-llm/tnt-llm.ipynb | 757 +++++++++++++++++++++++++++++++++ 1 file changed, 757 insertions(+) create mode 100644 examples/tnt-llm/tnt-llm.ipynb diff --git a/examples/tnt-llm/tnt-llm.ipynb b/examples/tnt-llm/tnt-llm.ipynb new file mode 100644 index 000000000..c8a9e2e84 --- /dev/null +++ b/examples/tnt-llm/tnt-llm.ipynb @@ -0,0 +1,757 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "86b44172-7131-44a3-a825-ac6a7347b7a7", + "metadata": {}, + "source": [ + "## Phase 1: Taxonomy Generation\n", + "\n", + "\n", + "Scratch: get runs" + ] + }, + { + "cell_type": "code", + "execution_count": 79, + "id": "a7e4f593-d3e5-4483-a2c7-b7e29c83a261", + "metadata": {}, + "outputs": [], + "source": [ + "# \"Data\"\n", + "\n", + "from datetime import datetime, timedelta\n", + "\n", + "from langsmith import Client\n", + "\n", + "client = Client(timeout_ms=30_000)\n", + "\n", + "yday = datetime.now() - timedelta(days=3)\n", + "runs = list(\n", + " client.list_runs(\n", + " # project_id=\"f53ccf51-57c1-4c97-afdf-7ca5569945cd\",\n", + " project_name=\"chat-langchain\",\n", + " filter=\"eq(is_root, true)\",\n", + " start_time=yday,\n", + " select=[\"inputs\", \"outputs\", \"feedback_stats\"],\n", + " )\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 111, + "id": "759c7ee9-ce44-42c8-8c6f-6a1e6bfe4f25", + "metadata": {}, + "outputs": [], + "source": [ + "import random\n", + "\n", + "\n", + "def run_to_doc(run) -> Doc:\n", + " turns = []\n", + " idx = 0\n", + " for turn in run.inputs[\"chat_history\"] or []:\n", + " key, value = next(iter(turn.items()))\n", + " turns.append(f\"<{key} idx={idx}>\\n{value}\\n\")\n", + " idx += 1\n", + " turns.append(\n", + " f\"\"\"\n", + "\n", + "{run.inputs['question']}\n", + "\"\"\"\n", + " )\n", + " if run.outputs and run.outputs[\"output\"]:\n", + " turns.append(\n", + " f\"\"\"\n", + "{run.outputs['output']}\n", + "\"\"\"\n", + " )\n", + " return {\n", + " \"id\": str(run.id),\n", + " \"content\": (\"\\n\".join(turns)),\n", + " }\n", + "\n", + "\n", + "docs = [run_to_doc(run) for run in runs if run.inputs]\n", + "docs = random.sample(docs, min(len(docs), 1000))" + ] + }, + { + "cell_type": "code", + "execution_count": 261, + "id": "cd9eda3e-3943-426a-bba2-8b9b81c00ec9", + "metadata": {}, + "outputs": [], + "source": [ + "use_case = (\n", + " \"Generate the taxonomy that can be used both to label the user intent\"\n", + " \" as well as to identify any required documentation (references, how-tos, etc.)\"\n", + " \" that would benefit the user.\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "a07c8526-614a-4154-abfe-91e768a58a46", + "metadata": {}, + "source": [ + "#### 1.a Summarize Docs" + ] + }, + { + "cell_type": "code", + "execution_count": 272, + "id": "ff02c2a1-18b5-4848-96bb-27ff00978570", + "metadata": {}, + "outputs": [], + "source": [ + "import operator\n", + "import re\n", + "from typing import Annotated, List, Optional, Sequence, TypedDict\n", + "\n", + "from langchain import hub\n", + "from langchain_anthropic import ChatAnthropic\n", + "from langchain_core.messages import AIMessage, BaseMessage, HumanMessage\n", + "from langchain_core.output_parsers import StrOutputParser\n", + "from langchain_core.prompts import ChatPromptTemplate\n", + "from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough\n", + "\n", + "\n", + "class Doc(TypedDict):\n", + " id: str\n", + " content: str\n", + " summary: Optional[str]\n", + " explanation: Optional[str]\n", + "\n", + "\n", + "class TaxonomyGenerationState(TypedDict):\n", + " documents: List[Doc]\n", + " # Indices\n", + " minibatches: List[List[int]]\n", + " # Candidate Taxonomies\n", + " clusters: Annotated[List[List[dict]], operator.add]\n", + "\n", + "\n", + "# Phase 1.a: Generate summaries + explanations\n", + "\n", + "\n", + "summary_prompt = hub.pull(\"wfh/tnt-llm-summary-generation\").partial(\n", + " summary_length=20, explanation_length=30\n", + ")\n", + "\n", + "\n", + "def parse_summary(xml_string: str) -> dict:\n", + " summary_pattern = r\"(.*?)\"\n", + " explanation_pattern = r\"(.*?)\"\n", + "\n", + " summary_match = re.search(summary_pattern, xml_string, re.DOTALL)\n", + " explanation_match = re.search(explanation_pattern, xml_string, re.DOTALL)\n", + "\n", + " summary = summary_match.group(1).strip() if summary_match else \"\"\n", + " explanation = explanation_match.group(1).strip() if explanation_match else \"\"\n", + "\n", + " return {\"summary\": summary, \"explanation\": explanation}\n", + "\n", + "\n", + "summary_llm_chain = (\n", + " summary_prompt\n", + " | ChatAnthropic(model=\"claude-3-haiku-20240307\")\n", + " | StrOutputParser()\n", + " # Customize the tracing name for easier organization\n", + ").with_config(run_name=\"GenerateSummary\")\n", + "summary_chain = summary_llm_chain | parse_xml\n", + "\n", + "\n", + "# Now combine as a \"map\" operation in a map-reduce chain\n", + "# Input: state\n", + "# Output: state U summaries\n", + "# Processes docs in parallel\n", + "def get_content(state: TaxonomyGenerationState):\n", + " docs = state[\"documents\"]\n", + " return [{\"content\": doc[\"content\"]} for doc in docs]\n", + "\n", + "\n", + "map_step = RunnablePassthrough.assign(\n", + " summaries=get_content\n", + " | RunnableLambda(func=summary_chain.batch, afunc=summary_chain.abatch)\n", + ")\n", + "\n", + "\n", + "def reduce_summaries(combined: dict) -> TaxonomyGenerationState:\n", + " summaries = combined[\"summaries\"]\n", + " documents = combined[\"documents\"]\n", + " return {\n", + " \"documents\": [\n", + " {\n", + " \"id\": doc[\"id\"],\n", + " \"content\": doc[\"content\"],\n", + " \"summary\": summ_info[\"summary\"],\n", + " \"explanation\": summ_info[\"explanation\"],\n", + " }\n", + " for doc, summ_info in zip(documents, summaries)\n", + " ]\n", + " }\n", + "\n", + "\n", + "# This is actually the node itself!\n", + "map_reduce_chain = map_step | reduce_summaries" + ] + }, + { + "cell_type": "code", + "execution_count": 114, + "id": "69c28ede-7a6c-485b-9606-f39e8ec215c3", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Failed to batch ingest runs: LangSmithConnectionError('Connection error caused failure to post https://api.smith.langchain.com/runs/batch in LangSmith API. Please confirm your internet connection.. SSLError(MaxRetryError(\"HTTPSConnectionPool(host=\\'api.smith.langchain.com\\', port=443): Max retries exceeded with url: /runs/batch (Caused by SSLError(SSLEOFError(8, \\'EOF occurred in violation of protocol (_ssl.c:2393)\\')))\"))')\n", + "Failed to batch ingest runs: LangSmithConnectionError('Connection error caused failure to post https://api.smith.langchain.com/runs/batch in LangSmith API. Please confirm your internet connection.. SSLError(MaxRetryError(\"HTTPSConnectionPool(host=\\'api.smith.langchain.com\\', port=443): Max retries exceeded with url: /runs/batch (Caused by SSLError(SSLEOFError(8, \\'EOF occurred in violation of protocol (_ssl.c:2393)\\')))\"))')\n" + ] + } + ], + "source": [ + "summarized_docs = map_reduce_chain.invoke({\"documents\": docs}, {\"max_concurrency\": 5})" + ] + }, + { + "cell_type": "code", + "execution_count": 121, + "id": "df55b666-0123-4fce-b15c-7126ff84acdb", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "\n", + "with open(\"docs.json\", \"w\") as f:\n", + " json.dump(summarized_docs, f)" + ] + }, + { + "cell_type": "markdown", + "id": "654cbad5-3af9-4c16-9d43-cbb903957944", + "metadata": {}, + "source": [ + "#### 1.b Split into Minibatches" + ] + }, + { + "cell_type": "code", + "execution_count": 236, + "id": "3e0139c3-b5ba-42b9-9367-33533d66eb58", + "metadata": {}, + "outputs": [], + "source": [ + "def get_minibatches(state: TaxonomyGenerationState, config: RunnableConfig):\n", + " batch_size = config[\"configurable\"].get(\"batch_size\", 200)\n", + " original = state[\"documents\"]\n", + " indices = list(range(len(original)))\n", + " random.shuffle(indices)\n", + " if len(indices) < batch_size:\n", + " # Don't pad needlessly if we can't fill a single batch\n", + " return [indices]\n", + "\n", + " num_full_batches = len(indices) // batch_size\n", + "\n", + " batches = [\n", + " indices[i * batch_size : (i + 1) * batch_size] for i in range(num_full_batches)\n", + " ]\n", + "\n", + " leftovers = len(indices) % batch_size\n", + " if leftovers:\n", + " last_batch = indices[num_full_batches * batch_size :]\n", + " elements_to_add = batch_size - leftovers\n", + " last_batch += random.sample(indices, elements_to_add)\n", + " batches.append(last_batch)\n", + "\n", + " return {\n", + " \"minibatches\": batches,\n", + " }" + ] + }, + { + "cell_type": "code", + "execution_count": 243, + "id": "870071b4-cba9-4711-8af3-70a6f7ecd225", + "metadata": {}, + "outputs": [], + "source": [ + "batched_state = get_minibatches(summarized_docs, {\"configurable\": {}})\n", + "batched_state = {**summarized_docs, **batched_state}" + ] + }, + { + "cell_type": "markdown", + "id": "1ca9ca18-f43a-46bd-bd28-7bec90932dd9", + "metadata": {}, + "source": [ + "#### 1.c Generate Taxonomy" + ] + }, + { + "cell_type": "code", + "execution_count": 273, + "id": "224ed013-2963-489c-b734-315cad701d59", + "metadata": {}, + "outputs": [], + "source": [ + "import random\n", + "import re\n", + "from typing import Any, Dict\n", + "\n", + "# We instruct our LLMs to generate 10\n", + "# intent categories and 25 domain categories for taxonomy generation.\n", + "\n", + "\n", + "def parse_taxa(output_text: str) -> Dict:\n", + " cluster_matches = re.findall(\n", + " r\"\\s*(.*?)\\s*(.*?)\\s*(.*?)\\s*\",\n", + " output_text,\n", + " re.DOTALL,\n", + " )\n", + " clusters = [\n", + " {\"id\": id.strip(), \"name\": name.strip(), \"description\": description.strip()}\n", + " for id, name, description in cluster_matches\n", + " ]\n", + " # We don't parse the explanation since it isn't used downstream\n", + " return {\"clusters\": clusters}\n", + "\n", + "\n", + "def format_docs(docs: List[Doc]) -> str:\n", + " xml_table = \"\\n\"\n", + " for doc in docs:\n", + " xml_table += f'{doc[\"summary\"]}\\n'\n", + " xml_table += \"\"\n", + " return xml_table\n", + "\n", + "\n", + "def format_taxonomy(clusters):\n", + " xml = \"\\n\"\n", + " for label in clusters:\n", + " xml += \" \\n\"\n", + " xml += f' {label[\"id\"]}\\n'\n", + " xml += f' {label[\"name\"]}\\n'\n", + " xml += f' {label[\"description\"]}\\n'\n", + " xml += \" \\n\"\n", + " xml += \"\"\n", + " return xml\n", + "\n", + "\n", + "# def generate_taxonomy(\n", + "# state: TaxonomyGenerationState, config: RunnableConfig\n", + "# ) -> TaxonomyGenerationState:\n", + "# \"\"\"Prompt an LLM to generate an initial taxonomy.\"\"\"\n", + "# configurable = config[\"configurable\"]\n", + "# docs = state[\"documents\"]\n", + "# mb_indices = state[\"minibatches\"][0]\n", + "# first_minibatch = [docs[idx] for idx in mb_indices]\n", + "# data_table_xml = format_docs(first_minibatch)\n", + "# initial_taxonomy = generate_taxonomy_chain.invoke(\n", + "# {\n", + "# \"data_xml\": data_table_xml,\n", + "# \"cluster_name_length\": configurable.get(\"cluster_name_length\", 10),\n", + "# \"cluster_description_length\": configurable.get(\n", + "# \"cluster_description_length\", 30\n", + "# ),\n", + "# \"explanation_length\": configurable.get(\"explanation_length\", 20),\n", + "# \"max_num_clusters\": configurable.get(\"max_num_clusters\", 25),\n", + "# }\n", + "# )\n", + "\n", + "# return {\n", + "# \"clusters\": [initial_taxonomy[\"clusters\"]],\n", + "# }" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c2f11188-3d1d-4d44-a5cf-886d46884b48", + "metadata": {}, + "outputs": [], + "source": [ + "taxonomy_generation_prompt = hub.pull(\"wfh/tnt-llm-taxonomy-generation\").partial(\n", + " use_case=\"Generate the taxonomy that can be used to label the user intent in the conversation.\",\n", + ")\n", + "\n", + "taxonomy_generation_llm = ChatAnthropic(\n", + " model=\"claude-3-haiku-20240307\", max_tokens_to_sample=2000\n", + ")\n", + "taxa_gen_llm_chain = (\n", + " taxonomy_generation_prompt | taxonomy_generation_llm | StrOutputParser()\n", + ").with_config(run_name=\"GenerateTaxonomy\")\n", + "\n", + "\n", + "generate_taxonomy_chain = taxa_gen_llm_chain | parse_generation_output" + ] + }, + { + "cell_type": "code", + "execution_count": 258, + "id": "f4b2820e-e002-4d51-8efe-feb32ab83bd1", + "metadata": {}, + "outputs": [], + "source": [ + "taxonomies = generate_taxonomy(batched_state, {\"configurable\": {\"max_concurrency\": 5}})\n", + "taxonomies = {**batched_state, **taxonomies}" + ] + }, + { + "cell_type": "markdown", + "id": "6df7a089-b268-47b6-b8a3-91606cfe1bec", + "metadata": {}, + "source": [ + "#### 1.c Update taxonomy\n", + "\n", + "Mostly the same." + ] + }, + { + "cell_type": "code", + "execution_count": 274, + "id": "ce10047e-7062-4935-8bd7-ffacba4bd099", + "metadata": {}, + "outputs": [], + "source": [ + "taxonomy_update_prompt = hub.pull(\"wfh/tnt-llm-taxonomy-update\")\n", + "\n", + "taxa_update_llm_chain = (\n", + " taxonomy_update_prompt | taxonomy_update_llm | StrOutputParser()\n", + ").with_config(run_name=\"UpdateTaxonomy\")\n", + "\n", + "\n", + "update_taxonomy_chain = taxa_update_llm_chain | parse_taxa\n", + "\n", + "# def update_taxonomy(\n", + "# state: TaxonomyGenerationState, config: RunnableConfig\n", + "# ) -> TaxonomyGenerationState:\n", + "# \"\"\"Prompt an LLM to update the taxonomy based on the current state.\"\"\"\n", + "# configurable = config[\"configurable\"]\n", + "# docs = state[\"documents\"]\n", + "# minibatches = state[\"minibatches\"]\n", + "# previous_taxonomy = state[\"clusters\"][-1]\n", + "# which_mb = len(state[\"clusters\"]) % len(minibatches)\n", + "# mb_indices = minibatches[which_mb]\n", + "# minibatch = [docs[idx] for idx in mb_indices]\n", + "# # The new data we will be using to\n", + "# data_table_xml = format_docs(minibatch)\n", + "# cluster_table_xml = format_taxonomy(previous_taxonomy)\n", + "# updated_taxonomy = update_taxonomy_chain.invoke(\n", + "# {\n", + "# \"data_xml\": data_table_xml,\n", + "# \"use_case\": configurable[\"use_case\"],\n", + "# \"cluster_table_xml\": cluster_table_xml,\n", + "# \"suggestion_length\": configurable.get(\"suggestion_length\", 30),\n", + "# \"cluster_name_length\": configurable.get(\"cluster_name_length\", 10),\n", + "# \"cluster_description_length\": configurable.get(\n", + "# \"cluster_description_length\", 30\n", + "# ),\n", + "# \"explanation_length\": configurable.get(\"explanation_length\", 20),\n", + "# \"max_num_clusters\": configurable.get(\"max_num_clusters\", 25),\n", + "# }\n", + "# )\n", + "# return {\n", + "# \"clusters\": [updated_taxonomy[\"clusters\"]],\n", + "# }" + ] + }, + { + "cell_type": "code", + "execution_count": 270, + "id": "6bf6b09f-5777-49df-bf47-5ae5ff3d2bce", + "metadata": {}, + "outputs": [], + "source": [ + "updated_taxonomies = update_taxonomy(\n", + " taxonomies, {\"configurable\": {\"use_case\": use_case}}\n", + ")\n", + "updated_taxonomies = {**taxonomies, **updated_taxonomies}" + ] + }, + { + "cell_type": "code", + "execution_count": 271, + "id": "e6c532c4-5468-4e35-a557-35ea7f0b5651", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[{'id': '1',\n", + " 'name': 'Manage file operations',\n", + " 'description': 'Load, read, write, and list files in various formats like JSON, CSV, and Parquet.'},\n", + " {'id': '2',\n", + " 'name': 'Resolve code and network issues',\n", + " 'description': 'Debug directory errors, network firewall problems, and create AI-powered models.'},\n", + " {'id': '3',\n", + " 'name': 'Implement conversational AI',\n", + " 'description': 'Handle user queries, provide responses, and manage user-AI interactions.'},\n", + " {'id': '4',\n", + " 'name': 'Build chatbots with LangChain',\n", + " 'description': 'Develop chatbots that use LangChain for knowledge retrieval, conversation history, and response generation.'},\n", + " {'id': '5',\n", + " 'name': 'Integrate LangChain with external services',\n", + " 'description': 'Combine LangChain with tools like Ollama, Mistral, and Azure OpenAI for advanced capabilities.'},\n", + " {'id': '6',\n", + " 'name': 'Process web content with LangChain',\n", + " 'description': 'Load web pages, extract text, and handle various file formats from online sources.'},\n", + " {'id': '7',\n", + " 'name': 'Manage LangChain vector stores',\n", + " 'description': 'Ingest, store, and retrieve text data using vector embeddings and database solutions.'}]" + ] + }, + "execution_count": 271, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "updated_taxonomies[\"clusters\"][-1]" + ] + }, + { + "cell_type": "markdown", + "id": "c2abc558-507f-44b7-9812-54fb1e9b3c64", + "metadata": {}, + "source": [ + "#### 1.d Review Taxonomy" + ] + }, + { + "cell_type": "code", + "execution_count": 275, + "id": "fe3ae6e6-80ac-4e71-8a43-5d8c0d64fa2c", + "metadata": {}, + "outputs": [], + "source": [ + "taxonomy_review_prompt = hub.pull(\"wfh/tnt-llm-taxonomy-review\")\n", + "\n", + "taxa_review_llm_chain = (\n", + " taxonomy_generation_llm | taxonomy_review_llm | StrOutputParser()\n", + ").with_config(run_name=\"ReviewTaxonomy\")\n", + "\n", + "\n", + "review_taxonomy_chain = taxa_review_llm_chain | parse_taxa\n", + "\n", + "\n", + "# def review_taxonomy(\n", + "# state: TaxonomyGenerationState, config: RunnableConfig\n", + "# ) -> TaxonomyGenerationState:\n", + "# \"\"\"Prompt an LLM to update the taxonomy based on the current state.\"\"\"\n", + "# configurable = config[\"configurable\"]\n", + "# docs = state[\"documents\"]\n", + "# minibatches = state[\"minibatches\"]\n", + "# previous_taxonomy = state[\"clusters\"][-1]\n", + "# mb_indices = random.sample\n", + "# which_mb = len(state[\"clusters\"]) % len(minibatches)\n", + "# mb_indices = minibatches[which_mb]\n", + "# minibatch = [docs[idx] for idx in mb_indices]\n", + "# # The new data we will be using to\n", + "# data_table_xml = format_docs(minibatch)\n", + "# cluster_table_xml = format_taxonomy(previous_taxonomy)\n", + "# reviewed_taxonomy = review_taxonomy_chain.invoke(\n", + "# {\n", + "# \"data_xml\": data_table_xml,\n", + "# \"use_case\": configurable[\"use_case\"],\n", + "# \"cluster_table_xml\": cluster_table_xml,\n", + "# \"suggestion_length\": configurable.get(\"suggestion_length\", 30),\n", + "# \"cluster_name_length\": configurable.get(\"cluster_name_length\", 10),\n", + "# \"cluster_description_length\": configurable.get(\n", + "# \"cluster_description_length\", 30\n", + "# ),\n", + "# \"explanation_length\": configurable.get(\"explanation_length\", 20),\n", + "# \"max_num_clusters\": configurable.get(\"max_num_clusters\", 25),\n", + "# }\n", + "# )\n", + "# return {\n", + "# \"clusters\": [reviewed_taxonomy[\"clusters\"]],\n", + "# }" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0039cf1c-54d5-4e9e-8dd6-a5cebfaec92d", + "metadata": {}, + "outputs": [], + "source": [ + "def invoke_taxonomy_chain(\n", + " chain: Callable,\n", + " state: TaxonomyGenerationState,\n", + " config: RunnableConfig,\n", + " mb_indices: List[int],\n", + ") -> TaxonomyGenerationState:\n", + " configurable = config[\"configurable\"]\n", + " docs = state[\"documents\"]\n", + " minibatch = [docs[idx] for idx in mb_indices]\n", + " data_table_xml = format_docs(minibatch)\n", + "\n", + " previous_taxonomy = state[\"clusters\"][-1] if state[\"clusters\"] else []\n", + " cluster_table_xml = format_taxonomy(previous_taxonomy)\n", + "\n", + " updated_taxonomy = chain.invoke(\n", + " {\n", + " \"data_xml\": data_table_xml,\n", + " \"use_case\": configurable[\"use_case\"],\n", + " \"cluster_table_xml\": cluster_table_xml,\n", + " \"suggestion_length\": configurable.get(\"suggestion_length\", 30),\n", + " \"cluster_name_length\": configurable.get(\"cluster_name_length\", 10),\n", + " \"cluster_description_length\": configurable.get(\n", + " \"cluster_description_length\", 30\n", + " ),\n", + " \"explanation_length\": configurable.get(\"explanation_length\", 20),\n", + " \"max_num_clusters\": configurable.get(\"max_num_clusters\", 25),\n", + " }\n", + " )\n", + "\n", + " return {\n", + " **state,\n", + " \"clusters\": state[\"clusters\"] + [updated_taxonomy[\"clusters\"]],\n", + " }\n", + "\n", + "\n", + "def generate_taxonomy(\n", + " state: TaxonomyGenerationState, config: RunnableConfig\n", + ") -> TaxonomyGenerationState:\n", + " return invoke_taxonomy_chain(\n", + " generate_taxonomy_chain, state, config, state[\"minibatches\"][0]\n", + " )\n", + "\n", + "\n", + "def update_taxonomy(\n", + " state: TaxonomyGenerationState, config: RunnableConfig\n", + ") -> TaxonomyGenerationState:\n", + " which_mb = len(state[\"clusters\"]) % len(state[\"minibatches\"])\n", + " return invoke_taxonomy_chain(\n", + " update_taxonomy_chain, state, config, state[\"minibatches\"][which_mb]\n", + " )\n", + "\n", + "\n", + "def review_taxonomy(\n", + " state: TaxonomyGenerationState, config: RunnableConfig\n", + ") -> TaxonomyGenerationState:\n", + " return invoke_taxonomy_chain(\n", + " review_taxonomy_chain, state, config, random.choice(state[\"minibatches\"])\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": 276, + "id": "e407115c-85bd-415f-8012-2ea061b1215c", + "metadata": {}, + "outputs": [], + "source": [ + "reviewed_taxonomies = review_taxonomy(\n", + " updated_taxonomies, {\"configurable\": {\"use_case\": use_case}}\n", + ")\n", + "reviewed_taxonomies = {**updated_taxonomies, **reviewed_taxonomies}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f1f97ea4-53e5-4f55-8d73-b5b2234a47d9", + "metadata": {}, + "outputs": [], + "source": [ + "# Define the graph\n", + "graph = StateGraph(TaxonomyGenerationState)\n", + "graph.add_node(\"summarize\", map_reduce_chain)\n", + "graph.add_node(\"get_minibatches\", get_minibatches)\n", + "graph.add_node(\"generate_taxonomy\", generate_taxonomy)\n", + "graph.add_node(\"update_taxonomy\", update_taxonomy)\n", + "graph.add_node(\"review_taxonomy\", review_taxonomy)\n", + "\n", + "graph.add_edge(\"summarize\", \"add_minibatches\")\n", + "graph.add_edge(\"add_minibatches\", \"generate_taxonomy\")\n", + "graph.add_edge(\"generate_taxonomy\", \"update_taxonomy\")\n", + "def should_review(\n", + "graph.add_conditional_edge(should_review)\n", + "graph.set_finish_point(\"review_taxonomy\")\n", + "\n", + "graph.set_entry_point(\"summarize\")\n", + "app = graph.compile()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9e1da39e-4c8f-407b-8011-f54931df7c7c", + "metadata": {}, + "outputs": [], + "source": [ + "# Phase 2: Text Classification\n", + "\n", + "\n", + "# Define the state for the text classification graph\n", + "class TextClassificationState(TypedDict):\n", + " messages: Annotated[Sequence[BaseMessage], operator.add]\n", + " labels: Annotated[Sequence[str], operator.add]\n", + "\n", + "\n", + "def assign_labels(state: TextClassificationState) -> TextClassificationState:\n", + " \"\"\"Prompt an LLM to assign labels to the given text.\"\"\"\n", + " llm = ChatOpenAI(temperature=0)\n", + " messages = [HumanMessage(content=text) for text in state[\"messages\"]]\n", + " responses = llm.batch(messages)\n", + " labels = [response.additional_kwargs[\"labels\"] for response in responses]\n", + " return {\"messages\": messages, \"labels\": labels}\n", + "\n", + "\n", + "# Define the TnT-LLM workflow using LangGraph\n", + "\n", + "# Phase 1: Taxonomy Generation\n", + "taxonomy_graph = StateGraph(TaxonomyGenerationState)\n", + "taxonomy_graph.add_node(\"summarize\", summarize_text)\n", + "taxonomy_graph.add_node(\"generate\", generate_taxonomy)\n", + "taxonomy_graph.add_node(\"update\", update_taxonomy)\n", + "taxonomy_graph.add_node(\"review\", review_taxonomy)\n", + "\n", + "taxonomy_graph.add_edge(\"summarize\", \"generate\")\n", + "taxonomy_graph.add_edge(\"generate\", \"update\")\n", + "taxonomy_graph.add_conditional_edges(\n", + " \"update\",\n", + " lambda state: \"update\" if len(state[\"summaries\"]) < 10 else \"review\",\n", + ")\n", + "taxonomy_graph.add_edge(\"review\", END)\n", + "\n", + "taxonomy_graph.set_entry_point(\"summarize\")\n", + "taxonomy_generator = taxonomy_graph.compile()\n", + "\n", + "# Phase 2: Text Classification\n", + "classification_graph = StateGraph(TextClassificationState)\n", + "classification_graph.add_node(\"assign_labels\", assign_labels)\n", + "\n", + "classification_graph.add_edge(\"assign_labels\", END)\n", + "\n", + "classification_graph.set_entry_point(\"assign_labels\")\n", + "text_classifier = classification_graph.compile()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.2" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}